from __future__ import annotations import logging from contextlib import nullcontext from dataclasses import dataclass from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.dp_attention import get_is_extend_in_batch from sglang.srt.layers.moe.token_dispatcher.base import ( BaseDispatcher, BaseDispatcherConfig, CombineInput, CombineInputFormat, DispatcherBaseHooks, DispatchOutput, DispatchOutputFormat, ) from sglang.srt.layers.moe.topk import TopKOutput from sglang.srt.layers.moe.utils import ( DeepEPMode, get_deepep_config, get_moe_runner_backend, is_tbo_enabled, ) from sglang.srt.utils import ( get_bool_env_var, is_blackwell, is_hip, is_npu, load_json_config, ) _is_npu = is_npu() if TYPE_CHECKING: from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs try: from deep_ep import Buffer, Config if not _is_npu: from sglang.srt.layers.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8, ) use_deepep = True except ImportError: use_deepep = False from enum import Enum, IntEnum, auto import torch import torch.distributed as dist _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() logger = logging.getLogger(__name__) class DeepEPPDispatchHooks(DispatcherBaseHooks): def __call__(self, dispatcher: BaseDispatcher): for hook_fun in self.hook_dict.values(): hook_fun(dispatcher) class DeepEPNormalDispatchOutput(NamedTuple): """DeepEP normal dispatch output.""" hidden_states: torch.Tensor hidden_states_scale: Optional[torch.Tensor] topk_ids: torch.Tensor topk_weights: torch.Tensor num_recv_tokens_per_expert: List[int] @property def format(self) -> DispatchOutputFormat: return DispatchOutputFormat.DEEPEP_NORMAL class DeepEPLLDispatchOutput(NamedTuple): """DeepEP low latency dispatch output.""" hidden_states: torch.Tensor hidden_states_scale: Optional[torch.Tensor] topk_ids: torch.Tensor topk_weights: torch.Tensor masked_m: torch.Tensor expected_m: int @property def format(self) -> DispatchOutputFormat: return DispatchOutputFormat.DEEPEP_LL assert isinstance(DeepEPNormalDispatchOutput, DispatchOutput) assert isinstance(DeepEPLLDispatchOutput, DispatchOutput) class DeepEPNormalCombineInput(NamedTuple): """DeepEP normal combine input.""" hidden_states: torch.Tensor topk_ids: torch.Tensor topk_weights: torch.Tensor @property def format(self) -> CombineInputFormat: return CombineInputFormat.DEEPEP_NORMAL class DeepEPLLCombineInput(NamedTuple): """DeepEP low latency combine input.""" hidden_states: torch.Tensor topk_ids: torch.Tensor topk_weights: torch.Tensor @property def format(self) -> CombineInputFormat: return CombineInputFormat.DEEPEP_LL assert isinstance(DeepEPNormalCombineInput, CombineInput) assert isinstance(DeepEPLLCombineInput, CombineInput) class DeepEPDispatchMode(IntEnum): NORMAL = auto() LOW_LATENCY = auto() class DeepEPBuffer: _buffer = None _dispatch_mode: Optional[DeepEPDispatchMode] = None _hidden_size: Optional[int] = None _num_max_dispatch_tokens_per_rank: Optional[int] = None _num_experts: Optional[int] = None @classmethod def get_deepep_buffer( cls, group: dist.ProcessGroup, hidden_size: int, param_bytes: int, deepep_mode: DeepEPMode, num_max_dispatch_tokens_per_rank: int = -1, num_experts: int = -1, ): if cls._buffer is not None: return cls._buffer cls._hidden_size = hidden_size cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank cls._num_experts = num_experts num_nvl_bytes, num_rdma_bytes = 0, 0 if deepep_mode.enable_normal(): hidden_bytes = hidden_size * param_bytes for config in ( DeepEPConfig.get_instance().normal_dispatch_config or Buffer.get_dispatch_config(group.size()), DeepEPConfig.get_instance().normal_combine_config or Buffer.get_combine_config(group.size()), ): num_nvl_bytes = max( config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes, ) num_rdma_bytes = max( config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes, ) if deepep_mode.enable_low_latency(): assert num_max_dispatch_tokens_per_rank != -1 assert num_experts != -1 and num_experts % group.size() == 0 num_rdma_bytes = max( Buffer.get_low_latency_rdma_size_hint( num_max_dispatch_tokens_per_rank, hidden_size, group.size(), num_experts, ), num_rdma_bytes, ) # We should calculate num_qps_per_rank consistently with DeepEP's test script logic: if deepep_mode == DeepEPMode.NORMAL: # refer: https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_internode.py#L235 num_qps_per_rank = DeepEPConfig.get_instance().num_sms elif deepep_mode == DeepEPMode.LOW_LATENCY: # refer: https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_low_latency.py#L176 num_qps_per_rank = num_experts // group.size() elif deepep_mode == DeepEPMode.AUTO: # low-latency and normal mode all need run # refer: https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_internode.py#L235 num_qps_per_rank = max( DeepEPConfig.get_instance().num_sms, num_experts // group.size() ) else: raise NotImplementedError if not _is_npu: total_num_sms = torch.cuda.get_device_properties( device="cuda" ).multi_processor_count if ( (deepep_mode != DeepEPMode.LOW_LATENCY) and not is_tbo_enabled() and (DeepEPConfig.get_instance().num_sms < total_num_sms // 2) ): logger.warning( f"Only use {DeepEPConfig.get_instance().num_sms} SMs for DeepEP communication. " f"This may result in highly suboptimal performance. " f"Consider using --deepep-config to change the behavior." ) cls._buffer = Buffer( group, num_nvl_bytes, num_rdma_bytes, low_latency_mode=deepep_mode.enable_low_latency(), num_qps_per_rank=num_qps_per_rank, # TODO can be false when unneeded allow_mnnvl=True, ) return cls._buffer @classmethod def clean_buffer(cls): if not cls._buffer.low_latency_mode: return cls._buffer.clean_low_latency_buffer( cls._num_max_dispatch_tokens_per_rank, cls._hidden_size, cls._num_experts, ) @classmethod def set_dispatch_mode_as_normal(cls): cls._dispatch_mode = DeepEPDispatchMode.NORMAL @classmethod def set_dispatch_mode_as_low_latency(cls): if cls._dispatch_mode == DeepEPDispatchMode.NORMAL: cls.clean_buffer() cls._dispatch_mode = DeepEPDispatchMode.LOW_LATENCY @classmethod def set_dispatch_mode(cls, mode: DeepEPMode): if mode.is_low_latency(): cls.set_dispatch_mode_as_low_latency() elif mode.is_normal(): cls.set_dispatch_mode_as_normal() else: raise Exception("unsupported mode") class DeepEPConfig(BaseDispatcherConfig): _instance = None def __init__(self): config_str = get_deepep_config() if config_str: config_parsed = load_json_config(config_str) if torch.distributed.get_rank() == 0: logger.info(f"Use DeepEP Config: {config_parsed}") config_dispatch = config_parsed["normal_dispatch"] config_combine = config_parsed["normal_combine"] self.normal_dispatch_config = Config(**config_dispatch) self.normal_combine_config = Config(**config_combine) assert config_dispatch["num_sms"] == config_combine["num_sms"] self.num_sms = config_dispatch["num_sms"] else: self.normal_dispatch_config = None self.normal_combine_config = None self.num_sms = Buffer.num_sms @classmethod def get_instance(cls): if cls._instance is None: cls._instance = DeepEPConfig() return cls._instance class _DeepEPDispatcherImplBase: def __init__( self, group: torch.distributed.ProcessGroup, router_topk: int, permute_fusion: bool, num_experts: int, num_local_experts: int, hidden_size: int, params_dtype: torch.dtype, deepep_mode: DeepEPMode, ): if not use_deepep: raise ImportError( "DeepEP is not installed. Please install DeepEP package from " "https://github.com/deepseek-ai/deepep." ) self.group = group self.router_topk = router_topk self.permute_fusion = permute_fusion self.num_experts = num_experts self.num_local_experts = num_local_experts self.hidden_size = hidden_size self.params_dtype = params_dtype self.deepep_mode = deepep_mode self.params_bytes = 2 # A large value will lead to large memory occupation, thus users should change it accordingly self.num_max_dispatch_tokens_per_rank = ( envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() ) # DeepEP internode_ll dispatch uses FINISHED_SUM_TAG=1024 # and the logic requires num-tokens-sent-from-one-rank-to-another-rank less than it assert self.num_max_dispatch_tokens_per_rank <= 1024 self.handle = None self.quant_config: Optional[dict] = None self.overlap_args: Optional[CombineOverlapArgs] = None self.meta_overlap_args: Optional[dict] = None def dispatch_a( self, hidden_states: torch.Tensor, topk_output: TopKOutput, ): raise NotImplementedError def dispatch_b(self, *args, **kwargs): raise NotImplementedError def combine_a( self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, ): raise NotImplementedError def combine_b(self, *args, **kwargs): raise NotImplementedError def _get_buffer(self): raise NotImplementedError def set_quant_config(self, quant_config: dict) -> None: self.quant_config = quant_config def set_overlap_args( self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict ) -> None: self.overlap_args = combine_overlap_args self.meta_overlap_args = meta_overlap_args def clear_overlap_args(self) -> None: self.overlap_args = None self.meta_overlap_args = None class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase): def __init__(self, async_finish: bool, **kwargs): super().__init__(**kwargs) self.async_finish = async_finish self.src2dst = None self.quant_config = {} def dispatch_a( self, hidden_states: torch.Tensor, topk_output: TopKOutput, ): topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids topk_ids = topk_ids.to(torch.int64) if ( deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and not get_moe_runner_backend().is_cutlass() and not envs.SGLANG_DEEPEP_BF16_DISPATCH.get() ): # TODO hard code 128 block quant,use fp8 communication hidden_states = sglang_per_token_group_quant_fp8( hidden_states, 128, column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, ) previous_event = Buffer.capture() if self.async_finish else None return hidden_states, topk_ids, topk_weights, previous_event def dispatch_b(self, hidden_states, topk_ids, topk_weights, previous_event): ( hidden_states, topk_ids, topk_weights, num_recv_tokens_per_expert, event, ) = self._dispatch_core(hidden_states, topk_ids, topk_weights, previous_event) event.current_stream_wait() if self.async_finish else () if isinstance(hidden_states, tuple): hidden_states, hidden_states_scale = hidden_states else: hidden_states_scale = None return DeepEPNormalDispatchOutput( hidden_states, hidden_states_scale, topk_ids, topk_weights, num_recv_tokens_per_expert, ) def _dispatch_core( self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], topk_ids: torch.Tensor, topk_weights: torch.Tensor, previous_event, ): buffer = self._get_buffer() ( num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event, ) = buffer.get_dispatch_layout( topk_ids, self.num_experts, previous_event=previous_event, async_finish=self.async_finish, allocate_on_comm_stream=previous_event is not None, ) # FIXME: `handle` should be transmitted with tokens from dispatch to combine. # However, doing this would incur an unknown synchronization error, but keeping # `handle` as a member variable works. ( recv_x, recv_topk_ids, recv_topk_weights, num_recv_tokens_per_expert, self.handle, event, ) = buffer.dispatch( x, topk_idx=topk_ids, topk_weights=topk_weights, num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, is_token_in_rank=is_token_in_rank, num_tokens_per_expert=num_tokens_per_expert, previous_event=previous_event, async_finish=self.async_finish, allocate_on_comm_stream=(previous_event is not None) and self.async_finish, expert_alignment=128 if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM else 1, config=DeepEPConfig.get_instance().normal_dispatch_config, ) get_global_expert_distribution_recorder().on_deepep_dispatch_normal( num_recv_tokens_per_expert, num_tokens_per_rank=num_tokens_per_rank, num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, num_tokens_per_expert=num_tokens_per_expert, ) return ( recv_x, recv_topk_ids, recv_topk_weights, num_recv_tokens_per_expert, event, ) def combine_a( self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, ): if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM or _use_aiter or _is_npu: output = hidden_states else: raise NotImplementedError() # triton runner was supported but it's temporarily disabled previous_event = Buffer.capture() if self.async_finish else None return output, previous_event def combine_b(self, output, previous_event): hidden_states, event = self._combine_core(output, previous_event) event.current_stream_wait() if self.async_finish else () self.handle = None self.src2dst = None return hidden_states def _combine_core(self, x: torch.Tensor, previous_event): buffer = self._get_buffer() combined_x, _, event = buffer.combine( x, self.handle, async_finish=self.async_finish, previous_event=previous_event, allocate_on_comm_stream=previous_event is not None, config=DeepEPConfig.get_instance().normal_combine_config, ) return combined_x, event def _get_buffer(self): DeepEPBuffer.set_dispatch_mode_as_normal() return DeepEPBuffer.get_deepep_buffer( self.group, self.hidden_size, self.params_bytes, self.deepep_mode, self.num_max_dispatch_tokens_per_rank, self.num_experts, ) class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): def __init__(self, return_recv_hook: bool, **kwargs): super().__init__(**kwargs) """ num_max_dispatch_tokens_per_rank: the actual batch size in the decoding engine should be less than 256 https://github.com/deepseek-ai/DeepEP?tab=readme-ov-file#example-use-in-inference-decoding """ self.return_recv_hook = return_recv_hook self.device_module = torch.get_device_module() self.quant_config = {} def dispatch_a( self, hidden_states: torch.Tensor, topk_output: TopKOutput, ): buffer = self._get_buffer() topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids topk_ids = topk_ids.to(torch.int64) expected_m = ( hidden_states.shape[0] * buffer.group_size * topk_ids.shape[1] + self.num_experts ) // self.num_experts hidden_states, masked_m, event, hook = self._dispatch_core( hidden_states, topk_ids, ) return ( hidden_states, topk_ids, topk_weights, masked_m, expected_m, event, hook, ) def dispatch_b( self, hidden_states, topk_ids, topk_weights, masked_m, expected_m, event, hook, ): hook() if self.return_recv_hook else event.current_stream_wait() get_global_expert_distribution_recorder().on_deepep_dispatch_low_latency( masked_m ) if isinstance(hidden_states, tuple): hidden_states, hidden_states_scale = hidden_states else: hidden_states_scale = None deepep_output = DeepEPLLDispatchOutput( hidden_states, hidden_states_scale, topk_ids, topk_weights, masked_m, expected_m, ) return deepep_output def _dispatch_core( self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, ): use_nvfp4 = use_fp8 = False input_global_scale = self.quant_config.get("input_global_scale", None) if input_global_scale is not None: use_nvfp4 = True elif not envs.SGLANG_DEEPEP_BF16_DISPATCH.get(): use_fp8 = True buffer = self._get_buffer() packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = ( buffer.low_latency_dispatch( hidden_states, topk_ids, self.num_max_dispatch_tokens_per_rank, self.num_experts, use_fp8=use_fp8, **(dict(use_nvfp4=True) if use_nvfp4 else dict()), **( dict(x_global_scale=input_global_scale) if input_global_scale is not None else dict() ), async_finish=not self.return_recv_hook, return_recv_hook=self.return_recv_hook, round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and deep_gemm_wrapper.DEEPGEMM_BLACKWELL, use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and deep_gemm_wrapper.DEEPGEMM_BLACKWELL, ) ) return packed_recv_hidden, self.packed_recv_count, event, hook def combine_a( self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, ): hidden_states, event, hook = self._combine_core( hidden_states, topk_ids, topk_weights, ) return hidden_states, event, hook def combine_b(self, hidden_states, event, hook): overlap_args = self.overlap_args if overlap_args is not None: overlap_args.stream.wait_stream(self.device_module.current_stream()) hook() if self.return_recv_hook else event.current_stream_wait() if overlap_args is not None: self.device_module.current_stream().wait_stream(overlap_args.stream) return hidden_states def _combine_core( self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, ): buffer = self._get_buffer() overlap_args = self.overlap_args meta_overlap_args = self.meta_overlap_args ctx = nullcontext() if overlap_args is not None: overlap_args.stream.wait_event(overlap_args.wait_event) ctx = torch.cuda.stream(overlap_args.stream) if is_blackwell(): overlap_args_dict = dict( overlap=overlap_args.overlap, src_signals=overlap_args.signal, src_signal_expect_value=overlap_args.threshold, ) else: overlap_args_dict = dict( overlap=overlap_args.overlap, packed_recv_count=self.packed_recv_count, comp_signal=overlap_args.signal, block_m=meta_overlap_args["block_m"], threshold=meta_overlap_args["threshold"], num_sms=overlap_args.num_sms, ) else: overlap_args_dict = {} with ctx: combined_hidden_states, event, hook = buffer.low_latency_combine( x=hidden_states, topk_idx=topk_ids, topk_weights=topk_weights, handle=self.handle, async_finish=not self.return_recv_hook, return_recv_hook=self.return_recv_hook, **overlap_args_dict, ) self.packed_recv_count = self.handle = None return combined_hidden_states, event, hook def _get_buffer(self): DeepEPBuffer.set_dispatch_mode_as_low_latency() return DeepEPBuffer.get_deepep_buffer( self.group, self.hidden_size, self.params_bytes, self.deepep_mode, self.num_max_dispatch_tokens_per_rank, self.num_experts, ) @dataclass class _Stage(Enum): INITIAL = auto() AFTER_DISPATCH_A = auto() AFTER_DISPATCH_B = auto() AFTER_COMBINE_A = auto() class DeepEPDispatcher(BaseDispatcher): def __init__( self, group: torch.distributed.ProcessGroup, router_topk: int, permute_fusion: bool = False, num_experts: int = None, num_local_experts: int = None, hidden_size: int = None, params_dtype: torch.dtype = None, deepep_mode: DeepEPMode = DeepEPMode.AUTO, async_finish: bool = False, return_recv_hook: bool = False, ): super().__init__() self.deepep_mode = deepep_mode common_kwargs = dict( group=group, router_topk=router_topk, permute_fusion=permute_fusion, num_experts=num_experts, num_local_experts=num_local_experts, hidden_size=hidden_size, params_dtype=params_dtype, deepep_mode=deepep_mode, ) if self.deepep_mode.enable_low_latency(): self._low_latency_dispatcher = _DeepEPDispatcherImplLowLatency( return_recv_hook=return_recv_hook, **common_kwargs, ) if self.deepep_mode.enable_normal(): self._normal_dispatcher = _DeepEPDispatcherImplNormal( async_finish=async_finish, **common_kwargs, ) self._stage = _Stage.INITIAL self._deepep_dispatch_hooks = DeepEPPDispatchHooks() def dispatch( self, hidden_states: torch.Tensor, topk_output: TopKOutput, ) -> DispatchOutput: self.dispatch_a(hidden_states, topk_output) if self._deepep_dispatch_hooks is not None: self._deepep_dispatch_hooks(self) ret = self.dispatch_b() return ret def dispatch_a( self, hidden_states: torch.Tensor, topk_output: TopKOutput, ): self._update_stage(_Stage.INITIAL, _Stage.AFTER_DISPATCH_A) inner_state = self._get_impl().dispatch_a( hidden_states=hidden_states, topk_output=topk_output, ) self._dispatch_intermediate_state = inner_state def dispatch_b(self): self._update_stage(_Stage.AFTER_DISPATCH_A, _Stage.AFTER_DISPATCH_B) inner_state = self._dispatch_intermediate_state del self._dispatch_intermediate_state return self._get_impl().dispatch_b(*inner_state) def combine( self, combine_input: CombineInput, ) -> torch.Tensor: self.combine_a(combine_input) ret = self.combine_b() return ret def combine_a( self, combine_input: CombineInput, ): hidden_states, topk_ids, topk_weights = combine_input self._update_stage(_Stage.AFTER_DISPATCH_B, _Stage.AFTER_COMBINE_A) inner_state = self._get_impl().combine_a( hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, ) self._combine_intermediate_state = inner_state def combine_b(self): self._update_stage(_Stage.AFTER_COMBINE_A, _Stage.INITIAL) inner_state = self._combine_intermediate_state del self._combine_intermediate_state return self._get_impl().combine_b(*inner_state) def _get_impl(self) -> _DeepEPDispatcherImplBase: is_extend_in_batch = get_is_extend_in_batch() resolved_deepep_mode = self.deepep_mode.resolve(is_extend_in_batch) if resolved_deepep_mode == DeepEPMode.NORMAL: return self._normal_dispatcher elif resolved_deepep_mode == DeepEPMode.LOW_LATENCY: return self._low_latency_dispatcher else: raise ValueError(f"Invalid deepep_mode: {self.deepep_mode}") def _update_stage(self, old_stage, new_stage): assert self._stage == old_stage self._stage = new_stage def set_quant_config(self, quant_config: dict): super().set_quant_config(quant_config) if self.deepep_mode.enable_low_latency(): self._low_latency_dispatcher.set_quant_config(quant_config) if self.deepep_mode.enable_normal(): self._normal_dispatcher.set_quant_config(quant_config) def set_overlap_args( self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict ): super().set_overlap_args(combine_overlap_args, meta_overlap_args) if self.deepep_mode.enable_low_latency(): self._low_latency_dispatcher.set_overlap_args( combine_overlap_args, meta_overlap_args ) if self.deepep_mode.enable_normal(): self._normal_dispatcher.set_overlap_args( combine_overlap_args, meta_overlap_args ) def clear_overlap_args(self): super().clear_overlap_args() if self.deepep_mode.enable_low_latency(): self._low_latency_dispatcher.clear_overlap_args() if self.deepep_mode.enable_normal(): self._normal_dispatcher.clear_overlap_args() def register_deepep_dispatch_hook(self, hook): return self._deepep_dispatch_hooks.register_hook(hook)