diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 93246893e..80719a3b0 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/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, diff --git a/python/sglang/srt/arg_groups/moe_hook.py b/python/sglang/srt/arg_groups/moe_hook.py index bb96534e3..09311778c 100644 --- a/python/sglang/srt/arg_groups/moe_hook.py +++ b/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. diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 8383132c0..21902574f 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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" diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 3b5de8c67..5e37f4aa4 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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) diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 5f352f156..fc54952d2 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/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( diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index ec684bbd1..df368bf91 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py index beb2dc372..6748059db 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep_v2.py +++ b/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( diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 16091ad96..1876e158d 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/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 diff --git a/test/manual/ep/test_deepep_v2_deterministic.py b/test/manual/ep/test_deepep_v2_deterministic.py new file mode 100644 index 000000000..e05443428 --- /dev/null +++ b/test/manual/ep/test_deepep_v2_deterministic.py @@ -0,0 +1,91 @@ +"""Check exact logprobs across batches against a running DeepEP v2 server. + +Example on four NVIDIA GPUs with DeepEP v2 installed (test BF16 and FP8 models): + + EP_DISABLE_GIN=1 NCCL_CUMEM_ENABLE=1 EP_REUSE_NCCL_COMM=0 \ + SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK=4096 \ + python -m sglang.launch_server --model-path Qwen/Qwen3-30B-A3B-FP8 \ + --tp-size 4 --ep-size 4 --moe-a2a-backend deepep_v2 \ + --deepep-v2-mode direct --attention-backend triton \ + --enable-deterministic-inference --disable-radix-cache \ + --chunked-prefill-size 4096 --cuda-graph-max-bs-decode 32 + + python test/manual/ep/test_deepep_v2_deterministic.py + +Radix caching must be disabled: otherwise repeated prompts can reuse the same +prefill and hide batch-dependent results. Keep DeepGEMM precompilation enabled. +""" + +import argparse + +import requests + + +def check_batch_invariance(base_url, batch_sizes, max_new_tokens): + response = requests.get(f"{base_url}/get_server_info", timeout=30) + response.raise_for_status() + info = response.json() + assert info["moe_a2a_backend"] == "deepep_v2", info["moe_a2a_backend"] + assert info["enable_deterministic_inference"], "Enable deterministic inference" + assert info["disable_radix_cache"], "Disable radix caching to test fresh prefill" + + prompts = [ + "Tell me about Richard Feynman: ", + "The capital city of France is", + "Explain why the sky is blue in three sentences.", + ] + for prompt in prompts: + reference = None + for mixed in (False, True): + for batch_size in batch_sizes: + texts = [prompt] * batch_size + if mixed: + # Put the target last, after unrelated variable-length prompts, + # to move its tokens between TP shards as the batch changes. + texts = [ + "Write a Python function that sorts a list. " * (i % 5 + 1) + for i in range(batch_size - 1) + ] + [prompt] + response = requests.post( + f"{base_url}/generate", + json={ + "text": texts, + "sampling_params": { + "temperature": 0, + "max_new_tokens": max_new_tokens, + "ignore_eos": True, + }, + "return_logprob": True, + }, + timeout=600, + ) + response.raise_for_status() + signatures = [ + [(p[0], p[1]) for p in r["meta_info"]["output_token_logprobs"]] + for r in response.json() + ] + assert len(signatures) == batch_size + target = signatures[-1] + assert len(target) == max_new_tokens + if reference is None: + reference = target + compared = [target] if mixed else signatures + assert all(s == reference for s in compared), ( + f"Logprobs differ: {prompt=!r}, {batch_size=}, {mixed=}" + ) + print(f"PASS {prompt=!r} {batch_size=} {mixed=}", flush=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:30000") + parser.add_argument( + "--batch-sizes", nargs="+", type=int, default=[1, 2, 3, 4, 8, 16, 32] + ) + parser.add_argument("--max-new-tokens", type=int, default=32) + args = parser.parse_args() + if args.max_new_tokens < 1 or any(bs < 1 for bs in args.batch_sizes): + parser.error("Token count and batch sizes must be positive") + check_batch_invariance( + args.base_url.rstrip("/"), args.batch_sizes, args.max_new_tokens + ) diff --git a/test/registered/kernel/communication/test_deterministic_reduce_scatter.py b/test/registered/kernel/communication/test_deterministic_reduce_scatter.py new file mode 100644 index 000000000..703718857 --- /dev/null +++ b/test/registered/kernel/communication/test_deterministic_reduce_scatter.py @@ -0,0 +1,80 @@ +"""A token must retain its bits when batch size changes its receiving TP rank. + +Run with ``python test_deterministic_reduce_scatter.py --num-gpu 4``. +This exercises the collective used before DeepEP MoE without a model or DeepEP. +""" + +import os + +import pytest +import torch +import torch.distributed as dist + +from sglang.srt.distributed import parallel_state as ps +from sglang.srt.environ import envs +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.utils import multigpu_pytest_main + +register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +@pytest.fixture(scope="module") +def group(): + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + # Match --enable-deterministic-inference's CUDA collective settings. + os.environ["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"] = "1" + os.environ["NCCL_ALGO"] = "allreduce:tree" + nchannels = str(envs.SGLANG_DETERMINISTIC_NCCL_NCHANNELS.get()) + os.environ["NCCL_MIN_NCHANNELS"] = nchannels + os.environ["NCCL_MAX_NCHANNELS"] = nchannels + torch.cuda.set_device(local_rank) + ps.set_custom_all_reduce(False) + ps.init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=local_rank, + distributed_init_method="env://", + ) + ps.initialize_model_parallel(tensor_model_parallel_size=world_size) + yield ps.get_tp_group() + ps.destroy_model_parallel() + ps.destroy_distributed_environment() + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("use_graph", [False, True]) +def test_batch_and_destination_invariance(group, dtype, use_graph): + torch.manual_seed(42 + group.rank_in_group) + # Non-integer contributions expose changes in floating-point sum order. + token = torch.randn(2048, device="cuda", dtype=dtype) + reference = None + for local_tokens in (1, 2, 3, 8, 32, 128): + input_ = token.repeat(local_tokens * group.world_size, 1) + original = input_.clone() + output = torch.empty((local_tokens, 2048), device="cuda", dtype=dtype) + if use_graph: + with group.graph_capture() as capture: + for _ in range(3): + group.reduce_scatter_tensor(output, input_) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture.stream): + group.reduce_scatter_tensor(output, input_) + graph.replay() + else: + group.reduce_scatter_tensor(output, input_) + + # Compare all destinations to catch rank-dependent reduction order. + all_outputs = torch.empty_like(input_) + dist.all_gather_into_tensor(all_outputs, output, group=group.device_group) + if reference is None: + reference = all_outputs[0].clone() + torch.testing.assert_close( + all_outputs, reference.expand_as(all_outputs), atol=0, rtol=0 + ) + torch.testing.assert_close(input_, original, atol=0, rtol=0) + + +if __name__ == "__main__": + multigpu_pytest_main(__name__, __file__, num_gpus=(4,)) diff --git a/test/registered/kernel/moe/test_deepep_v2_contig_scatter.py b/test/registered/kernel/moe/test_deepep_v2_contig_scatter.py new file mode 100644 index 000000000..ce201ee6d --- /dev/null +++ b/test/registered/kernel/moe/test_deepep_v2_contig_scatter.py @@ -0,0 +1,115 @@ +"""Tests for the DeepEP v2 contiguous-layout scatter kernel.""" + +import unittest + +import torch + +from sglang.kernels.ops.moe.ep_moe_kernels import ep_scatter_from_psum +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +DEVICE = "cuda" +HIDDEN = 256 +SCALE_HIDDEN = HIDDEN // 128 +# Match the scatter kernel's BLOCK_E alignment. +ALIGN = 128 + + +class TestDeepEPv2ContigScatter(CustomTestCase): + """BF16 and FP8 scatter must preserve routed rows and mask invalid experts.""" + + # Local expert ids per (token, slot); -1 marks a route to a remote expert, + # and 7 is out of this rank's expert range. + RECV_TOPK = [ + [0, -1], + [0, -1], + [0, 1], + [1, 7], + [-1, -1], + ] + NUM_LOCAL_EXPERTS = 2 + + def _run(self, dtype, with_scale): + num_recv = len(self.RECV_TOPK) + recv_x = ( + torch.arange(num_recv * HIDDEN, dtype=torch.float32, device=DEVICE).reshape( + num_recv, HIDDEN + ) + % 100 + ).to(dtype) + recv_topk = torch.tensor(self.RECV_TOPK, dtype=torch.int64, device=DEVICE) + psum = torch.tensor( + [ALIGN * (e + 1) for e in range(self.NUM_LOCAL_EXPERTS)], + dtype=torch.int32, + device=DEVICE, + ) + all_tokens = int(psum[-1].item()) + + recv_x_scale = None + output_tensor_scale = None + if with_scale: + recv_x_scale = torch.arange( + num_recv * SCALE_HIDDEN, dtype=torch.float32, device=DEVICE + ).reshape(num_recv, SCALE_HIDDEN) + output_tensor_scale = torch.zeros( + (all_tokens, SCALE_HIDDEN), dtype=torch.float32, device=DEVICE + ) + + output_tensor = torch.zeros((all_tokens, HIDDEN), device=DEVICE, dtype=dtype) + m_indices = torch.empty(all_tokens, device=DEVICE, dtype=torch.int32) + output_index = torch.empty_like(recv_topk) + expert_start_loc = torch.empty_like(psum) + + ep_scatter_from_psum( + recv_x, + recv_x_scale, + recv_topk, + psum, + expert_start_loc, + output_tensor, + output_tensor_scale, + m_indices, + output_index, + ) + return recv_x, recv_x_scale, output_tensor, output_tensor_scale, output_index + + def _check(self, dtype, with_scale): + recv_x, recv_x_scale, out, out_scale, output_index = self._run( + dtype, with_scale + ) + index = output_index.tolist() + + for token, slots in enumerate(self.RECV_TOPK): + for slot, expert in enumerate(slots): + dest = index[token][slot] + if not 0 <= expert < self.NUM_LOCAL_EXPERTS: + # -1 suppresses this route in the post-permute gather. + self.assertEqual(dest, -1, msg=f"{token=} {slot=} {expert=}") + continue + self.assertTrue( + ALIGN * expert <= dest < ALIGN * (expert + 1), + msg=f"{token=} {slot=} {expert=} {dest=}", + ) + torch.testing.assert_close(out[dest].float(), recv_x[token].float()) + if with_scale: + torch.testing.assert_close(out_scale[dest], recv_x_scale[token]) + + accepted = [ + index[t][s] + for t, slots in enumerate(self.RECV_TOPK) + for s, e in enumerate(slots) + if 0 <= e < self.NUM_LOCAL_EXPERTS + ] + self.assertEqual(len(set(accepted)), len(accepted)) + + def test_scatter_bf16_without_scales(self): + self._check(torch.bfloat16, with_scale=False) + + def test_scatter_fp8_with_scales(self): + self._check(torch.float8_e4m3fn, with_scale=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/distributed/test_parallel_state.py b/test/registered/unit/distributed/test_parallel_state.py index aad9a823f..f870b19e0 100644 --- a/test/registered/unit/distributed/test_parallel_state.py +++ b/test/registered/unit/distributed/test_parallel_state.py @@ -41,6 +41,7 @@ from contextlib import nullcontext from unittest.mock import Mock, patch import pytest +import torch from sglang.test.ci.ci_register import register_cpu_ci @@ -50,6 +51,87 @@ register_cpu_ci(est_time=11, suite="base-a-test-cpu") parallel_state = pytest.importorskip("sglang.srt.distributed.parallel_state") +@pytest.mark.parametrize("rank", range(4)) +@pytest.mark.parametrize("inplace_allreduce", [False, True]) +@pytest.mark.parametrize("alias_output", [False, True]) +def test_deterministic_reduce_scatter_preserves_input_and_selects_rank_shard( + monkeypatch, rank, inplace_allreduce, alias_output +): + monkeypatch.setenv("SGLANG_ENABLE_DETERMINISTIC_INFERENCE", "1") + coordinator = parallel_state.GroupCoordinator.__new__( + parallel_state.GroupCoordinator + ) + coordinator.rank_in_group = rank + coordinator.world_size = 4 + # Flattened input is a valid reduce-scatter layout too. + input_ = torch.arange(24, dtype=torch.float32) + original = input_.clone() + output = input_.view(4, 2, 3)[rank] if alias_output else torch.empty((2, 3)) + reduced = original + 100 + + def all_reduce(tensor): + assert tensor.data_ptr() != input_.data_ptr() + torch.testing.assert_close(tensor, original) + if inplace_allreduce: + tensor.copy_(reduced) + return tensor + return reduced + + coordinator.all_reduce = Mock(side_effect=all_reduce) + coordinator._reduce_scatter_tensor = Mock(side_effect=AssertionError) + coordinator.reduce_scatter_tensor(output, input_) + torch.testing.assert_close(output, reduced.view(4, 2, 3)[rank]) + if alias_output: + original.view(4, 2, 3)[rank].copy_(output) + torch.testing.assert_close(input_, original) + coordinator.all_reduce.assert_called_once() + + +@pytest.mark.parametrize("rank", range(4)) +@pytest.mark.parametrize("sizes", [None, [1, 3, 2, 0]]) +def test_deterministic_reduce_scatterv_selects_rank_shard(monkeypatch, rank, sizes): + monkeypatch.setenv("SGLANG_ENABLE_DETERMINISTIC_INFERENCE", "1") + coordinator = parallel_state.GroupCoordinator.__new__( + parallel_state.GroupCoordinator + ) + coordinator.rank_in_group = rank + coordinator.world_size = 4 + rows = 8 if sizes is None else sum(sizes) + input_ = torch.arange(rows * 3, dtype=torch.float32).view(rows, 3) + original = input_.clone() + reduced = original + 100 + + def all_reduce(tensor): + assert tensor.data_ptr() != input_.data_ptr() + torch.testing.assert_close(tensor, original) + return reduced + + coordinator.all_reduce = Mock(side_effect=all_reduce) + # pynccl must not be touched on the deterministic path. + coordinator.pynccl_comm = None + output = coordinator.reduce_scatterv(input_, sizes=sizes) + + offset = (rows // 4) * rank if sizes is None else sum(sizes[:rank]) + chunk = rows // 4 if sizes is None else sizes[rank] + torch.testing.assert_close(output, reduced.narrow(0, offset, chunk)) + torch.testing.assert_close(input_, original) + coordinator.all_reduce.assert_called_once() + + +def test_nondeterministic_reduce_scatter_keeps_native_path(monkeypatch): + monkeypatch.setenv("SGLANG_ENABLE_DETERMINISTIC_INFERENCE", "0") + monkeypatch.setattr(parallel_state, "_is_cpu", True) + coordinator = parallel_state.GroupCoordinator.__new__( + parallel_state.GroupCoordinator + ) + coordinator._reduce_scatter_tensor = Mock() + coordinator.all_reduce = Mock(side_effect=AssertionError) + input_ = torch.arange(8, dtype=torch.float32) + output = torch.empty(2) + coordinator.reduce_scatter_tensor(output, input_) + coordinator._reduce_scatter_tensor.assert_called_once_with(output, input_) + + def test_custom_allreduce_precedes_symmetric_memory_pynccl(): coordinator = parallel_state.GroupCoordinator.__new__( parallel_state.GroupCoordinator diff --git a/test/registered/unit/layers/moe/test_deepep_v2_wire_dtype.py b/test/registered/unit/layers/moe/test_deepep_v2_wire_dtype.py new file mode 100644 index 000000000..1da77e058 --- /dev/null +++ b/test/registered/unit/layers/moe/test_deepep_v2_wire_dtype.py @@ -0,0 +1,193 @@ +"""CPU-only tests for the DeepEP v2 BF16 / FP8 wire format.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.layers.moe.token_dispatcher import deepep_v2 +from sglang.srt.layers.moe.utils import ( + DeepEPv2Fp8ScaleFormat, + DispatcherOutputDtype, + get_deepep_v2_dispatcher_output_dtype, +) +from sglang.srt.runtime_context import get_context, get_exec, reset_context +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=4, suite="base-a-test-cpu") + +HIDDEN = 256 +TOPK = 4 +NUM_EXPERTS = 8 +NUM_LOCAL_EXPERTS = 2 +NUM_MAX_TOKENS = 16 + + +class _FakeGroup: + pass + + +class _FakeHandle: + def __init__(self, num_recv): + self.psum_num_recv_tokens_per_scaleup_rank = torch.tensor([num_recv]) + self.psum_num_recv_tokens_per_expert = torch.tensor( + [num_recv] * NUM_LOCAL_EXPERTS, dtype=torch.int32 + ) + + +class _FakeBuffer: + """Echoes the dispatch input back so the wire format stays observable.""" + + last = None + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + self.num_bytes = 1 << 20 + type(self).last = self + + def dispatch(self, x, **kwargs): + self.dispatch_x = x + num_recv = (x[0] if isinstance(x, tuple) else x).shape[0] + topk_idx = kwargs["topk_idx"] + topk_weights = kwargs["topk_weights"] + event = SimpleNamespace(event=None, current_stream_wait=lambda: None) + return x, topk_idx, topk_weights, _FakeHandle(num_recv), event + + +def _fake_quant(hidden_states, block_size, **kwargs): + return ( + torch.zeros_like(hidden_states, dtype=torch.float8_e4m3fn), + torch.zeros( + (hidden_states.shape[0], hidden_states.shape[1] // block_size), + dtype=torch.float32, + ), + ) + + +class _DeepEPv2WireDtypeBase(CustomTestCase): + def setUp(self): + reset_context() + self._published = get_context().override_server_args(model_path="dummy") + self._published.install() + _FakeBuffer.last = None + self._patches = [ + patch.object(deepep_v2, "use_deepep_v2", True), + patch.object(deepep_v2, "ElasticBuffer", _FakeBuffer, create=True), + patch.object(deepep_v2, "sglang_per_token_group_quant_fp8", _fake_quant), + patch.object(deepep_v2.dist, "get_world_size", return_value=4), + patch.object(deepep_v2.dist, "get_rank", return_value=0), + patch.object( + deepep_v2, + "get_deepep_v2_fp8_scale_format", + lambda: DeepEPv2Fp8ScaleFormat(tma_aligned=False, ue8m0=False), + ), + ] + for item in self._patches: + item.start() + + def tearDown(self): + self._published.restore() + reset_context() + for item in reversed(self._patches): + item.stop() + + def _dispatch(self, use_fp8_dispatch, num_tokens=8, is_extend_in_batch=True): + dispatcher = deepep_v2.DeepEPv2Dispatcher( + group=_FakeGroup(), + router_topk=TOPK, + num_experts=NUM_EXPERTS, + num_local_experts=NUM_LOCAL_EXPERTS, + hidden_size=HIDDEN, + params_dtype=torch.bfloat16, + use_fp8_dispatch=use_fp8_dispatch, + ) + dispatcher._impl.num_max_dispatch_tokens_per_rank = NUM_MAX_TOKENS + hidden_states = torch.randn((num_tokens, HIDDEN), dtype=torch.bfloat16) + topk_output = SimpleNamespace( + topk_ids=torch.zeros((num_tokens, TOPK), dtype=torch.int32), + topk_weights=torch.ones((num_tokens, TOPK), dtype=torch.float32), + ) + with patch.object( + deepep_v2, "get_is_extend_in_batch", lambda: is_extend_in_batch + ): + return hidden_states, dispatcher._impl.dispatch(hidden_states, topk_output) + + +class TestDeepEPv2WireDtype(_DeepEPv2WireDtypeBase): + def test_bf16_dispatch_sends_unquantized_activations(self): + hidden_states, out = self._dispatch(use_fp8_dispatch=False) + self.assertIs(_FakeBuffer.last.dispatch_x, hidden_states) + self.assertIsNone(out.hidden_states_scale) + self.assertEqual(out.hidden_states.dtype, torch.bfloat16) + self.assertFalse(out.hidden_states_scale_tma_aligned) + + def test_fp8_dispatch_still_sends_activations_and_scales(self): + _, out = self._dispatch(use_fp8_dispatch=True) + self.assertIsInstance(_FakeBuffer.last.dispatch_x, tuple) + self.assertIsNotNone(out.hidden_states_scale) + self.assertEqual(out.hidden_states.dtype, torch.float8_e4m3fn) + + def test_bf16_masked_decode_dispatch_stays_unquantized(self): + hidden_states, out = self._dispatch( + use_fp8_dispatch=False, is_extend_in_batch=False + ) + self.assertIs(_FakeBuffer.last.dispatch_x, hidden_states) + self.assertTrue(out.use_masked_gemm) + self.assertIsNone(out.hidden_states_scale) + + def test_wire_dtype_selects_the_elastic_buffer_layout(self): + for use_fp8_dispatch in (True, False): + with self.subTest(use_fp8_dispatch=use_fp8_dispatch): + _FakeBuffer.last = None + deepep_v2.DeepEPv2Buffer.destroy() + self._dispatch(use_fp8_dispatch=use_fp8_dispatch) + self.assertEqual( + _FakeBuffer.last.kwargs["use_fp8_dispatch"], use_fp8_dispatch + ) + + +class TestDeepEPv2DispatcherOutputDtype(CustomTestCase): + def setUp(self): + reset_context() + self._published = get_context().override_server_args(model_path="dummy") + self._published.install() + + def tearDown(self): + self._published.restore() + reset_context() + + def _flag(self, value): + return get_exec().moe.override(deepep_dispatcher_output_dtype=value) + + def test_auto_follows_the_checkpoint(self): + with self._flag("auto"): + self.assertIs( + get_deepep_v2_dispatcher_output_dtype(True), DispatcherOutputDtype.FP8 + ) + self.assertIs( + get_deepep_v2_dispatcher_output_dtype(False), DispatcherOutputDtype.BF16 + ) + + def test_explicit_flag_matching_the_checkpoint_is_accepted(self): + with self._flag("bf16"): + self.assertIs( + get_deepep_v2_dispatcher_output_dtype(False), DispatcherOutputDtype.BF16 + ) + with self._flag("fp8"): + self.assertIs( + get_deepep_v2_dispatcher_output_dtype(True), DispatcherOutputDtype.FP8 + ) + + def test_explicit_flag_contradicting_the_checkpoint_is_rejected(self): + with self._flag("bf16"): + with self.assertRaisesRegex(ValueError, "contradicts this checkpoint"): + get_deepep_v2_dispatcher_output_dtype(True) + with self._flag("fp8"): + with self.assertRaisesRegex(ValueError, "contradicts this checkpoint"): + get_deepep_v2_dispatcher_output_dtype(False) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py b/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py index e89b939da..2a3cb4ea7 100644 --- a/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py +++ b/test/registered/unit/layers/moe/test_hpc_ops_runner_guard.py @@ -106,17 +106,30 @@ def test_deepep_v2_quant_contract_rejects_incompatible_fp8( _validate_deepep_v2_quant_method(_fp8_method(**overrides)) +def test_deepep_v2_quant_contract_accepts_unquantized_bf16(_moe_flags): + """BF16 experts take the BF16 wire format instead of the FP8 one.""" + from sglang.srt.layers.moe.fused_moe_triton.layer import ( + _deepep_v2_experts_are_fp8, + _validate_deepep_v2_quant_method, + ) + from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod + + _moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2 + method = object.__new__(UnquantizedFusedMoEMethod) + _validate_deepep_v2_quant_method(method) + assert _deepep_v2_experts_are_fp8(method) is False + assert _deepep_v2_experts_are_fp8(_fp8_method()) is True + + def test_deepep_v2_quant_contract_rejects_incompatible_methods(_moe_flags): from sglang.srt.layers.moe.fused_moe_triton.layer import ( _validate_deepep_v2_quant_method, ) - from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod from sglang.srt.layers.quantization.w4afp8 import W4AFp8MoEMethod _moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2 - for method_type in (UnquantizedFusedMoEMethod, W4AFp8MoEMethod): - with pytest.raises(ValueError, match=method_type.__name__): - _validate_deepep_v2_quant_method(object.__new__(method_type)) + with pytest.raises(ValueError, match=W4AFp8MoEMethod.__name__): + _validate_deepep_v2_quant_method(object.__new__(W4AFp8MoEMethod)) def test_deepep_v2_quant_contract_does_not_affect_other_backends(_moe_flags): diff --git a/test/registered/unit/layers/moe/test_moe_runner_extensions.py b/test/registered/unit/layers/moe/test_moe_runner_extensions.py index 0452dd8a1..05969b2d7 100644 --- a/test/registered/unit/layers/moe/test_moe_runner_extensions.py +++ b/test/registered/unit/layers/moe/test_moe_runner_extensions.py @@ -147,7 +147,7 @@ def test_fused_moe_uses_explicit_quant_method_for_full_lifecycle(monkeypatch) -> monkeypatch.setattr( fused_moe_layer_module, "create_moe_dispatcher", - lambda config: SimpleNamespace(), + lambda config, quant_method: SimpleNamespace(), ) with ( @@ -346,7 +346,7 @@ def test_fused_moe_layer_runner_is_none_when_method_builds_no_runner( monkeypatch.setattr( fused_moe_layer_module, "create_moe_dispatcher", - lambda config: SimpleNamespace(), + lambda config, quant_method: SimpleNamespace(), ) with ( diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 891f831c5..d1f2d9236 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -2537,15 +2537,14 @@ class TestDeepEPv2Args(CustomTestCase): with self.assertRaisesRegex(ValueError, "instance connector"): handle_a2a_moe(args) - def test_deterministic_inference_rejected(self): + def test_deterministic_inference_accepted(self): args = self._args( moe_runner_backend="deep_gemm", enable_deterministic_inference=True, ) - with self.assertRaisesRegex(ValueError, "deterministic sorting"): - handle_a2a_moe(args) + handle_a2a_moe(args) - def test_rl_on_policy_deterministic_inference_rejected(self): + def test_rl_on_policy_deterministic_inference_accepted(self): args = self._args( moe_runner_backend="deep_gemm", rl_on_policy_target="fsdp", @@ -2558,8 +2557,7 @@ class TestDeepEPv2Args(CustomTestCase): ), ): handle_deterministic_inference(args) - with self.assertRaisesRegex(ValueError, "deterministic sorting"): - handle_a2a_moe(args) + handle_a2a_moe(args) def test_deterministic_inference_does_not_affect_legacy_deepep(self): args = self._args(