[Feature] Support BF16 and batch-invariant inference with DeepEP v2 (#38160)

python/sglang/kernels/ops/moe/ep_moe_kernels.py

@@ -1296,6 +1296,10 @@ def ep_scatter_from_psum(
         BLOCK_E=BLOCK_E,
     )
 
+    # The BF16 specialization never dereferences these scale pointers.
+    recv_x_scale_arg = recv_x_scale if is_fp8 else recv_x
+    output_tensor_scale_arg = output_tensor_scale if is_fp8 else output_tensor
+
     grid = min(recv_topk.shape[0], 1024 * 8)
     _fwd_kernel_ep_scatter_2[(grid,)](
         recv_topk.shape[0],
@@ -1303,7 +1307,7 @@ def ep_scatter_from_psum(
         recv_x,
         recv_x.stride(0),
         recv_x.stride(1),
-        recv_x_scale,
+        recv_x_scale_arg,
         recv_x_scale.stride(0) if is_fp8 else 0,
         recv_x_scale.stride(1) if is_fp8 else 0,
         recv_topk,
@@ -1312,12 +1316,15 @@ def ep_scatter_from_psum(
         output_tensor,
         output_tensor.stride(0),
         output_tensor.stride(1),
-        output_tensor_scale,
+        output_tensor_scale_arg,
         output_tensor_scale.stride(0) if is_fp8 else 0,
         output_tensor_scale.stride(1) if is_fp8 else 0,
         output_index,
         output_index.stride(0),
         output_index.stride(1),
+        # DeepEP v2 already rebases recv_topk to local expert IDs.
+        0,
+        num_experts,
         topk_num=recv_topk.shape[1],
         num_warps=num_warps,
         HIDDEN_SIZE=hidden_size,

python/sglang/srt/arg_groups/moe_hook.py

@@ -325,13 +325,6 @@ def handle_a2a_moe(server_args: Any):
 
     if a2a_backend == "deepep_v2":
         validate_deepep_v2_model_architecture(server_args)
-        if resolved_view(server_args).enable_deterministic_inference:
-            raise ValueError(
-                "DeepEP v2 does not forward deterministic=True to "
-                "ElasticBuffer, so deterministic sorting remains disabled. "
-                "Disable --enable-deterministic-inference or use "
-                "--moe-a2a-backend deepep."
-            )
         # ElasticBuffer requires CUMEM, but not NVLS or its preallocation.
         os.environ.setdefault("NCCL_CUMEM_ENABLE", "1")
         # Respect model-level runner declarations before resolving auto.

python/sglang/srt/distributed/parallel_state.py

@@ -1092,7 +1092,17 @@ class GroupCoordinator:
         return output
 
     def reduce_scatter_tensor(self, output: torch.Tensor, input: torch.Tensor):
-        if _is_npu or _is_cpu:
+        if envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get():
+            assert input.numel() == output.numel() * self.world_size
+            # Reduction order must be independent of the receiving rank.
+            # Preserve input even when all_reduce mutates its argument.
+            reduced = self.all_reduce(input.clone())
+            output.copy_(
+                reduced.reshape(-1)
+                .narrow(0, self.rank_in_group * output.numel(), output.numel())
+                .view_as(output)
+            )
+        elif _is_npu or _is_cpu:
             # TODO: add optimized reduce_scatter_tensor kernel for cpu
             self._reduce_scatter_tensor(output, input)
         elif self._maybe_aiter_reduce_scatter(output, input):
@@ -1173,6 +1183,33 @@ class GroupCoordinator:
         torch.distributed.reduce_scatter(output, input_list, group=self.device_group)
         return output
 
+    def _deterministic_reduce_scatterv(
+        self,
+        input_: torch.Tensor,
+        output: Optional[torch.Tensor],
+        sizes: Optional[List[int]],
+    ) -> torch.Tensor:
+        # Reduction order must be independent of the receiving rank.
+        # Offsets are a prefix sum, so unequal `sizes` work unchanged.
+        if sizes is not None:
+            assert len(sizes) == self.world_size
+            assert input_.shape[0] == sum(sizes)
+            chunk_size = sizes[self.rank_in_group]
+            offset = sum(sizes[: self.rank_in_group])
+        else:
+            assert input_.shape[0] % self.world_size == 0
+            chunk_size = input_.shape[0] // self.world_size
+            offset = chunk_size * self.rank_in_group
+        output_shape = (chunk_size,) + input_.shape[1:]
+        if output is None:
+            output = torch.empty(output_shape, dtype=input_.dtype, device=input_.device)
+        else:
+            assert output.shape == output_shape
+        # Preserve input even when all_reduce mutates its argument.
+        reduced = self.all_reduce(input_.clone())
+        output.copy_(reduced.narrow(0, offset, chunk_size))
+        return output
+
     def reduce_scatterv(
         self,
         input_: torch.Tensor,
@@ -1182,6 +1219,9 @@ class GroupCoordinator:
         world_size = self.world_size
         pynccl_comm = self.pynccl_comm
 
+        if envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get():
+            return self._deterministic_reduce_scatterv(input_, output, sizes)
+
         with pynccl_comm.change_state(enable=True):
             assert pynccl_comm is not None and not pynccl_comm.disabled, (
                 "pynccl is required for reduce_scatterv"

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -51,7 +51,9 @@ from sglang.srt.layers.moe.topk import (
     TopKOutputChecker,
 )
 from sglang.srt.layers.moe.utils import (
+    DispatcherOutputDtype,
     RoutingMethodType,
+    get_deepep_v2_dispatcher_output_dtype,
     has_per_rank_fused_shared_slots,
     uses_per_rank_fused_shared_slots,
 )
@@ -157,7 +159,10 @@ def _get_deepep_comm_group(a2a_backend):
     return group
 
 
-def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
+def create_moe_dispatcher(
+    moe_runner_config: MoeRunnerConfig,
+    quant_method: FusedMoEMethodBase,
+) -> BaseDispatcher:
     a2a_backend = get_moe_a2a_backend()
     if a2a_backend.is_none() and is_npu():
         return AscendTPDispatcher(moe_runner_config)
@@ -193,6 +198,9 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
             return_recv_hook=True,
         )
     elif a2a_backend.is_deepep_v2():
+        output_dtype = get_deepep_v2_dispatcher_output_dtype(
+            _deepep_v2_experts_are_fp8(quant_method)
+        )
         return DeepEPv2Dispatcher(
             group=get_tp_group().device_group,
             router_topk=moe_runner_config.top_k,
@@ -200,6 +208,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
             num_local_experts=moe_runner_config.num_local_experts,
             hidden_size=moe_runner_config.hidden_size,
             params_dtype=moe_runner_config.params_dtype,
+            use_fp8_dispatch=output_dtype is DispatcherOutputDtype.FP8,
         )
     elif a2a_backend.is_flashinfer():
         return FlashinferDispatcher(
@@ -239,11 +248,19 @@ def _validate_hpc_ops_quant_method(quant_method) -> None:
         )
 
 
+def _deepep_v2_experts_are_fp8(quant_method) -> bool:
+    # All other supported quantization methods are blockwise FP8.
+    return not isinstance(quant_method, UnquantizedFusedMoEMethod)
+
+
 def _validate_deepep_v2_quant_method(quant_method) -> None:
-    """Validate the FP8 contract consumed by the DeepEP v2 adapter."""
+    """Validate the expert formats the DeepEP v2 adapter can feed."""
     if not get_moe_a2a_backend().is_deepep_v2():
         return
 
+    if isinstance(quant_method, UnquantizedFusedMoEMethod):
+        return
+
     config = (
         quant_method.quant_config if isinstance(quant_method, Fp8MoEMethod) else None
     )
@@ -261,9 +278,10 @@ def _validate_deepep_v2_quant_method(quant_method) -> None:
 
     if reason is not None:
         raise ValueError(
-            "--moe-a2a-backend deepep_v2 requires 128x128 blockwise FP8 "
-            f"experts with dynamic activation scaling, but this layer {reason}. "
-            "Use a compatible checkpoint or --moe-a2a-backend deepep."
+            "--moe-a2a-backend deepep_v2 requires either 128x128 blockwise FP8 "
+            "experts with dynamic activation scaling or unquantized BF16 "
+            f"experts, but this layer {reason}. Use a compatible checkpoint or "
+            "--moe-a2a-backend deepep."
         )
 
 
@@ -499,7 +517,9 @@ class FusedMoE(torch.nn.Module):
         )
 
         self.quant_method.create_moe_runner(self, self.moe_runner_config)
-        self.dispatcher = create_moe_dispatcher(self.moe_runner_config)
+        self.dispatcher = create_moe_dispatcher(
+            self.moe_runner_config, quant_method=self.quant_method
+        )
         # Dispatchers are not nn.Modules, so they cannot register their own
         # buffers; the AITER expert mask would not survive a memory-saver resume.
         expert_mask = getattr(self.dispatcher, "expert_mask_gpu", None)

python/sglang/srt/layers/moe/moe_runner/deep_gemm.py

@@ -1692,10 +1692,12 @@ def pre_permute_deepep_v2_to_deep_gemm(
     deepep_v2_masked_max_m = dispatch_output.masked_max_m
     deepep_v2_total_expanded = dispatch_output.total_expanded
     deepep_v2_expert_alignment = dispatch_output.expert_alignment
-    if hidden_states_scale is None:
+    is_fp8 = hidden_states_scale is not None
+    if not is_fp8 and hidden_states.dtype != torch.bfloat16:
         raise RuntimeError(
-            "DeepEP v2 -> DeepGEMM requires FP8 dispatch output with activation "
-            "scales, but the dispatch output carried none."
+            "DeepEP v2 -> DeepGEMM requires either FP8 dispatch output with "
+            "activation scales or BF16 dispatch output, but the dispatch "
+            f"output carried {hidden_states.dtype} without scales."
         )
     assert runner_config.activation == "silu"
 
@@ -1764,7 +1766,9 @@ def pre_permute_deepep_v2_to_deep_gemm(
     input_tensor = torch.empty(
         (all_tokens, K), device=hidden_states.device, dtype=hidden_states.dtype
     )
-    if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+    if not is_fp8:
+        input_tensor_scale = None
+    elif deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
         # Packed UE8M0 scales require zero padding lanes.
         input_tensor_scale = torch.zeros(
             (ceil_div(K // 128, 4), all_tokens),
@@ -1792,7 +1796,8 @@ def pre_permute_deepep_v2_to_deep_gemm(
         scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
     )
     dispose_tensor(hidden_states)
-    dispose_tensor(hidden_states_scale)
+    if hidden_states_scale is not None:
+        dispose_tensor(hidden_states_scale)
     running_state["output_index"] = output_index
 
     return DeepGemmRunnerInput(

python/sglang/srt/layers/moe/moe_runner/runner.py

@@ -80,9 +80,11 @@ class MoeRunner:
             raise ValueError(
                 "--moe-a2a-backend deepep_v2 requires the deep_gemm MoE runner, "
                 f"but this MoE layer's quantization method selected the "
-                f"'{runner_backend.value}' runner. deepep_v2 dispatches FP8 "
-                "activations plus scales, which only deep_gemm consumes; use an "
-                "FP8 blockwise-quantized checkpoint, or --moe-a2a-backend deepep."
+                f"'{runner_backend.value}' runner. deepep_v2 dispatches into "
+                "the deep_gemm grouped-GEMM layout (FP8 activations plus "
+                "scales, or BF16 activations for unquantized experts); use an "
+                "FP8 blockwise-quantized or BF16 checkpoint, or "
+                "--moe-a2a-backend deepep."
             )
 
         self.fused_func = None

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

@@ -226,6 +226,7 @@ class _DeepEPv2Impl:
         hidden_size: int,
         scale_format: DeepEPv2Fp8ScaleFormat,
         num_max_dispatch_tokens_per_rank: int,
+        use_fp8_dispatch: bool,
     ):
         self.group = group
         self.router_topk = router_topk
@@ -234,6 +235,7 @@ class _DeepEPv2Impl:
         self.hidden_size = hidden_size
         self.scale_format = scale_format
         self.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank
+        self.use_fp8_dispatch = use_fp8_dispatch
         self.rank = dist.get_rank(group)
         self._handle = None
         self._pad_empty_combine = False
@@ -247,7 +249,7 @@ class _DeepEPv2Impl:
             self.hidden_size,
             self.router_topk,
             self.num_max_dispatch_tokens_per_rank,
-            True,
+            self.use_fp8_dispatch,
         )
 
     def _validate_common(
@@ -267,7 +269,7 @@ class _DeepEPv2Impl:
             )
         if self.hidden_size % _SCALE_BLOCK_SIZE != 0:
             raise ValueError(
-                "DeepEP v2 FP8 dispatch requires hidden_size multiple of "
+                "DeepEP v2 requires hidden_size multiple of "
                 f"{_SCALE_BLOCK_SIZE}, got {self.hidden_size}"
             )
         if topk_ids.shape[1] != self.router_topk:
@@ -302,8 +304,11 @@ class _DeepEPv2Impl:
             ).unsqueeze(0)
             topk_weights = topk_weights.new_zeros((1, topk_weights.shape[-1]))
 
-        _ensure_fp8_quant_available()
-        if use_masked:
+        if not self.use_fp8_dispatch:
+            dispatch_x = hidden_states
+            use_tma_aligned_col_major_sf = False
+        elif use_masked:
+            _ensure_fp8_quant_available()
             _ue8m0 = self.scale_format.ue8m0
             dispatch_x = sglang_per_token_group_quant_fp8(
                 hidden_states,
@@ -427,6 +432,7 @@ class DeepEPv2Dispatcher(BaseDispatcher):
         num_local_experts: int,
         hidden_size: int,
         params_dtype: torch.dtype,
+        use_fp8_dispatch: bool,
     ):
         super().__init__()
         if params_dtype != torch.bfloat16:
@@ -438,6 +444,7 @@ class DeepEPv2Dispatcher(BaseDispatcher):
         self.num_max_dispatch_tokens_per_rank = (
             envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
         )
+        self.use_fp8_dispatch = use_fp8_dispatch
         self._impl = _DeepEPv2Impl(
             group=group,
             router_topk=router_topk,
@@ -446,6 +453,7 @@ class DeepEPv2Dispatcher(BaseDispatcher):
             hidden_size=hidden_size,
             scale_format=scale_format,
             num_max_dispatch_tokens_per_rank=self.num_max_dispatch_tokens_per_rank,
+            use_fp8_dispatch=use_fp8_dispatch,
         )
 
     def dispatch(

python/sglang/srt/layers/moe/utils.py

@@ -405,6 +405,25 @@ def get_ascend_dispatcher_output_dtype(dispatcher):
     return DispatcherOutputDtype.BF16
 
 
+def get_deepep_v2_dispatcher_output_dtype(
+    experts_are_fp8: bool,
+) -> DispatcherOutputDtype:
+    """Match the dispatch dtype to the expert weights consumed by DeepGEMM."""
+    required = (
+        DispatcherOutputDtype.FP8 if experts_are_fp8 else DispatcherOutputDtype.BF16
+    )
+    requested = get_exec().moe.deepep_dispatcher_output_dtype
+    if requested != "auto" and DispatcherOutputDtype(requested) is not required:
+        raise ValueError(
+            f"--deepep-dispatcher-output-dtype {requested} contradicts this "
+            f"checkpoint: --moe-a2a-backend deepep_v2 dispatches "
+            f"{required.value} for "
+            f"{'FP8 blockwise' if experts_are_fp8 else 'BF16'} experts. Drop "
+            "the flag to let it follow the checkpoint."
+        )
+    return required
+
+
 def get_deepep_v2_fp8_scale_format() -> DeepEPv2Fp8ScaleFormat:
     """Resolve the FP8 scale layout DeepEP v2 must pre-quantize into."""
     from sglang.srt.layers import deep_gemm_wrapper