[Feature] Support BF16 and batch-invariant inference with DeepEP v2 (#38160)
This commit is contained in:
Cheng Wan
@@ -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