[Feature] Support BF16 and batch-invariant inference with DeepEP v2 (#38160)
This commit is contained in:
Cheng Wan
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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,))
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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):
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user