[Feature] Add DeepEPv2 (ElasticBuffer) MoE A2A backend (#35634)
Co-authored-by: menyu <menyu@nvidia.com> Co-authored-by: Jinyan Chen <93358689+liz-badada@users.noreply.github.com> Co-authored-by: Han Yu <helloyu0903@gmail.com> Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
MengYu
co-authored by
menyu
Jinyan Chen
Han Yu
Cheng Wan
menyu
Jinyan Chen
Han Yu
Cheng Wan
parent
cbfe54fba8
commit
a3ae667d67
@@ -0,0 +1,194 @@
|
||||
"""DP>1 routed-expert readback parity for DeepEP-family A2A backends."""
|
||||
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import pybase64
|
||||
import requests
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-h100")
|
||||
|
||||
_MODEL = os.environ.get("SGLANG_ROUTED_EXPERTS_TEST_MODEL", "deepseek-ai/DeepSeek-V3")
|
||||
_NUM_EXPERTS = 24
|
||||
_NUM_LAYERS = 1
|
||||
_TOPK = 8
|
||||
|
||||
_DUMMY_WEIGHT_ENV = {
|
||||
"SGLANG_ENABLE_ASYNC_ASSERT": "0",
|
||||
"SGLANG_SANITIZE_NAN_LOGITS": "1",
|
||||
"SGLANG_CUDA_COREDUMP": "0",
|
||||
"CUDA_ENABLE_COREDUMP_ON_EXCEPTION": "0",
|
||||
"SGLANG_CUDA_COREDUMP_BEFORE_CRASH": "0",
|
||||
}
|
||||
|
||||
|
||||
def _deep_ep_has(attr: str) -> bool:
|
||||
try:
|
||||
import deep_ep # noqa: F401
|
||||
except ImportError:
|
||||
return False
|
||||
return hasattr(deep_ep, attr)
|
||||
|
||||
|
||||
def _deep_ep_nccl_compatible() -> bool:
|
||||
try:
|
||||
version = torch.cuda.nccl.version()
|
||||
except (AttributeError, RuntimeError):
|
||||
return False
|
||||
return version is not None and version >= (2, 30, 7)
|
||||
|
||||
|
||||
class _ReadbackMixin:
|
||||
backend_args: list
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--json-model-override-args",
|
||||
json.dumps(
|
||||
{
|
||||
"num_hidden_layers": _NUM_LAYERS,
|
||||
"first_k_dense_replace": 0,
|
||||
"n_routed_experts": _NUM_EXPERTS,
|
||||
}
|
||||
),
|
||||
"--tp",
|
||||
"2",
|
||||
"--dp",
|
||||
"2",
|
||||
"--ep",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--enable-return-routed-experts",
|
||||
"--disable-cuda-graph",
|
||||
"--disable-radix-cache",
|
||||
# Keep the startup budget within the test's 256-token buffer.
|
||||
"--chunked-prefill-size",
|
||||
"256",
|
||||
"--mem-fraction-static",
|
||||
"0.5",
|
||||
*cls.backend_args,
|
||||
]
|
||||
cls.process = popen_launch_server(
|
||||
_MODEL,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
env={
|
||||
**os.environ,
|
||||
**_DUMMY_WEIGHT_ENV,
|
||||
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
|
||||
"SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256",
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
if getattr(cls, "process", None):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def _one_request(self, i: int):
|
||||
resp = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": f"{self._WORDS[i]} is item number {i}. Describe it in detail.",
|
||||
"sampling_params": {"max_new_tokens": 24, "temperature": 0},
|
||||
"return_routed_experts": True,
|
||||
},
|
||||
timeout=300,
|
||||
)
|
||||
self.assertEqual(resp.status_code, 200)
|
||||
meta = resp.json()["meta_info"]
|
||||
self.assertIn("routed_experts", meta)
|
||||
arr = np.frombuffer(pybase64.b64decode(meta["routed_experts"]), dtype=np.int32)
|
||||
self.assertEqual(
|
||||
arr.size % (_NUM_LAYERS * _TOPK),
|
||||
0,
|
||||
f"req{i}: payload size {arr.size} not a multiple of layers*topk",
|
||||
)
|
||||
rows = arr.reshape(-1, _NUM_LAYERS, _TOPK)
|
||||
self.assertGreater(rows.shape[0], 0)
|
||||
self.assertTrue(
|
||||
bool(((rows >= 0) & (rows < _NUM_EXPERTS)).all()),
|
||||
f"req{i}: expert id out of range [{rows.min()}, {rows.max()}]",
|
||||
)
|
||||
return rows
|
||||
|
||||
_WORDS = ["Alpha", "Bravo", "Charlie", "Delta", "Echo", "Foxtrot"]
|
||||
_N_REQ = 6
|
||||
|
||||
def test_dp2_readback(self):
|
||||
solo = [self._one_request(i) for i in range(self._N_REQ)]
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=self._N_REQ) as ex:
|
||||
conc = list(ex.map(self._one_request, range(self._N_REQ)))
|
||||
|
||||
for i in range(self._N_REQ):
|
||||
a, b = solo[i], conc[i]
|
||||
n = min(a.shape[0], b.shape[0])
|
||||
total = match = 0
|
||||
for t in range(n):
|
||||
for layer in range(_NUM_LAYERS):
|
||||
total += 1
|
||||
if set(a[t, layer].tolist()) == set(b[t, layer].tolist()):
|
||||
match += 1
|
||||
frac = match / max(1, total)
|
||||
self.assertGreaterEqual(
|
||||
frac,
|
||||
0.9,
|
||||
f"req{i}: only {frac:.1%} of per-token expert sets match the "
|
||||
"solo baseline — the capturer is reading rows that belong to "
|
||||
"other tokens (DeepEP-class backend misclassification)",
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_deep_ep_has("Buffer"), "DeepEP (v1 Buffer) not installed")
|
||||
class TestRoutedExpertsReadbackDeepEP(_ReadbackMixin, CustomTestCase):
|
||||
backend_args = [
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
"--deepep-mode",
|
||||
"low_latency",
|
||||
"--deepep-dispatcher-output-dtype",
|
||||
"fp8",
|
||||
"--moe-runner-backend",
|
||||
"deep_gemm",
|
||||
]
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_deep_ep_has("ElasticBuffer"), "DeepEP v2 (ElasticBuffer) not installed"
|
||||
)
|
||||
@unittest.skipUnless(
|
||||
_deep_ep_nccl_compatible(), "DeepEP v2 requires NCCL runtime >= 2.30.7"
|
||||
)
|
||||
class TestRoutedExpertsReadbackDeepEPv2(_ReadbackMixin, CustomTestCase):
|
||||
backend_args = [
|
||||
"--moe-a2a-backend",
|
||||
"deepep_v2",
|
||||
"--deepep-v2-mode",
|
||||
"direct",
|
||||
"--moe-runner-backend",
|
||||
"deep_gemm",
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,154 @@
|
||||
"""CPU-only tests for the DeepEP v2 ElasticBuffer ownership facade."""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.token_dispatcher import deepep_v2
|
||||
from sglang.srt.runtime_context import get_resources, reset_context
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class _FakeGroup:
|
||||
pass
|
||||
|
||||
|
||||
class _FakeBuffer:
|
||||
instances = []
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
self.num_bytes = 1 << 20
|
||||
type(self).instances.append(self)
|
||||
|
||||
|
||||
class TestDeepEPv2BufferLifecycle(CustomTestCase):
|
||||
def setUp(self):
|
||||
reset_context()
|
||||
_FakeBuffer.instances = []
|
||||
self._patches = [
|
||||
patch.object(deepep_v2, "use_deepep_v2", True),
|
||||
patch.object(deepep_v2, "ElasticBuffer", _FakeBuffer, create=True),
|
||||
patch.object(deepep_v2.dist, "get_world_size", return_value=8),
|
||||
]
|
||||
for item in self._patches:
|
||||
item.start()
|
||||
|
||||
def tearDown(self):
|
||||
reset_context()
|
||||
for item in reversed(self._patches):
|
||||
item.stop()
|
||||
|
||||
def _get(self, group=None, **overrides):
|
||||
kwargs = {
|
||||
"group": group or _FakeGroup(),
|
||||
"hidden_size": 4096,
|
||||
"router_topk": 8,
|
||||
"num_max_dispatch_tokens_per_rank": 128,
|
||||
"use_fp8_dispatch": True,
|
||||
"allow_hybrid_mode": False,
|
||||
}
|
||||
kwargs.update(overrides)
|
||||
return deepep_v2.DeepEPv2Buffer.get_buffer(**kwargs)
|
||||
|
||||
def test_same_key_reuses_buffer(self):
|
||||
group = _FakeGroup()
|
||||
first = self._get(group)
|
||||
second = self._get(group)
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(len(_FakeBuffer.instances), 1)
|
||||
|
||||
def test_constructor_inputs_participate_in_key(self):
|
||||
group = _FakeGroup()
|
||||
first = self._get(group)
|
||||
second = self._get(group, num_max_dispatch_tokens_per_rank=256)
|
||||
third = self._get(
|
||||
group,
|
||||
num_max_dispatch_tokens_per_rank=256,
|
||||
allow_hybrid_mode=True,
|
||||
)
|
||||
self.assertIsNot(first, second)
|
||||
self.assertIsNot(second, third)
|
||||
self.assertEqual(len(_FakeBuffer.instances), 3)
|
||||
|
||||
def test_key_keeps_process_group_object(self):
|
||||
group = _FakeGroup()
|
||||
self._get(group)
|
||||
state = get_resources().buffers[deepep_v2.DeepEPv2Buffer._STATE_KEY]
|
||||
self.assertIs(state.key[0], group)
|
||||
|
||||
def test_distinct_process_group_rebuilds(self):
|
||||
first = self._get(_FakeGroup())
|
||||
second = self._get(_FakeGroup())
|
||||
self.assertIsNot(first, second)
|
||||
self.assertEqual(len(_FakeBuffer.instances), 2)
|
||||
|
||||
def test_state_lives_in_runtime_resources(self):
|
||||
self._get()
|
||||
self.assertIn(
|
||||
deepep_v2.DeepEPv2Buffer._STATE_KEY,
|
||||
get_resources().buffers,
|
||||
)
|
||||
|
||||
def test_reset_context_drops_state_and_rebuilds(self):
|
||||
group = _FakeGroup()
|
||||
self._get(group)
|
||||
reset_context()
|
||||
self.assertNotIn(
|
||||
deepep_v2.DeepEPv2Buffer._STATE_KEY,
|
||||
get_resources().buffers,
|
||||
)
|
||||
self._get(group)
|
||||
self.assertEqual(len(_FakeBuffer.instances), 2)
|
||||
|
||||
def test_failed_constructor_is_not_published(self):
|
||||
class _FailingBuffer:
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise RuntimeError("construct failed")
|
||||
|
||||
with patch.object(deepep_v2, "ElasticBuffer", _FailingBuffer):
|
||||
with self.assertRaisesRegex(RuntimeError, "construct failed"):
|
||||
self._get()
|
||||
|
||||
state = get_resources().buffers[deepep_v2.DeepEPv2Buffer._STATE_KEY]
|
||||
self.assertIsNone(state.buffer)
|
||||
self.assertIsNone(state.key)
|
||||
self._get()
|
||||
self.assertEqual(len(_FakeBuffer.instances), 1)
|
||||
|
||||
def test_destroy_clears_facade_state(self):
|
||||
group = _FakeGroup()
|
||||
first = self._get(group)
|
||||
deepep_v2.DeepEPv2Buffer.destroy()
|
||||
state = get_resources().buffers[deepep_v2.DeepEPv2Buffer._STATE_KEY]
|
||||
self.assertIsNone(state.buffer)
|
||||
self.assertIsNone(state.key)
|
||||
second = self._get(group)
|
||||
self.assertIsNot(first, second)
|
||||
|
||||
def test_unavailable_deepep_fails_before_state_creation(self):
|
||||
with patch.object(deepep_v2, "use_deepep_v2", False):
|
||||
with self.assertRaisesRegex(ImportError, "github.com/deepseek-ai/DeepEP"):
|
||||
self._get()
|
||||
self.assertNotIn(
|
||||
deepep_v2.DeepEPv2Buffer._STATE_KEY,
|
||||
get_resources().buffers,
|
||||
)
|
||||
|
||||
def test_dispatch_capacity_guard_uses_actual_input_rows(self):
|
||||
impl = object.__new__(deepep_v2._DeepEPv2Impl)
|
||||
impl.num_max_dispatch_tokens_per_rank = 4
|
||||
impl.hidden_size = 128
|
||||
impl.router_topk = 2
|
||||
impl._validate_common(torch.empty(4, 128), torch.zeros(4, 2))
|
||||
with self.assertRaisesRegex(ValueError, "per-rank buffer capacity"):
|
||||
impl._validate_common(torch.empty(5, 128), torch.zeros(5, 2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,244 @@
|
||||
"""Tests for the DeepEP v2 expanded/masked repack kernels."""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.moe.ep_moe_kernels import (
|
||||
expand_to_masked_slab,
|
||||
masked_slab_to_expand,
|
||||
)
|
||||
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", runner_config="1-gpu-large")
|
||||
|
||||
DEVICE = "cuda"
|
||||
|
||||
|
||||
def _build_layout(counts, align, hidden, dtype, with_scale=False, scale_hidden=4):
|
||||
"""Build synthetic expanded-layout buffers for per-expert counts."""
|
||||
starts, psum = [], []
|
||||
prev_end = 0
|
||||
for c in counts:
|
||||
start = ((prev_end + align - 1) // align) * align
|
||||
end = start + c
|
||||
starts.append(start)
|
||||
psum.append(end)
|
||||
prev_end = end
|
||||
total = max(((prev_end + align - 1) // align) * align, 1)
|
||||
|
||||
# Vary rows and columns to expose broadcast or stride errors.
|
||||
base = torch.zeros((total, hidden), dtype=torch.float32, device=DEVICE)
|
||||
col_gain = 1.0 + (torch.arange(hidden, device=DEVICE) % 2).float()
|
||||
for s, c in zip(starts, counts):
|
||||
for j in range(c):
|
||||
base[s + j] = float((s + j) % 200 + 1) * col_gain
|
||||
recv_x = base.to(dtype)
|
||||
|
||||
scale = None
|
||||
if with_scale:
|
||||
scale = torch.zeros((total, scale_hidden), dtype=torch.float32, device=DEVICE)
|
||||
# Vary scale columns to expose pack-dimension stride errors.
|
||||
col = torch.arange(scale_hidden, dtype=torch.float32, device=DEVICE)
|
||||
for s, c in zip(starts, counts):
|
||||
for j in range(c):
|
||||
scale[s + j] = float((s + j) % 50 + 1) * 0.5 + col
|
||||
|
||||
psum_t = torch.tensor(psum, dtype=torch.int32, device=DEVICE)
|
||||
return recv_x, scale, psum_t, starts, total
|
||||
|
||||
|
||||
def _real_rows(starts, counts):
|
||||
rows = []
|
||||
for s, c in zip(starts, counts):
|
||||
rows.extend(range(s, s + c))
|
||||
return rows
|
||||
|
||||
|
||||
class TestDeepEPv2MaskedSlab(CustomTestCase):
|
||||
ALIGN = 16
|
||||
HIDDEN = 8
|
||||
MAX_M = 32
|
||||
|
||||
def _check_expand_roundtrip(self, counts, dtype, with_scale, topk=False):
|
||||
recv_x, scale, psum, starts, total = _build_layout(
|
||||
counts, self.ALIGN, self.HIDDEN, dtype, with_scale=with_scale
|
||||
)
|
||||
E = len(counts)
|
||||
masked_x, masked_x_scale, masked_m = expand_to_masked_slab(
|
||||
recv_x, scale, psum, E, self.MAX_M, self.ALIGN
|
||||
)
|
||||
|
||||
self.assertEqual(masked_m.tolist(), list(counts))
|
||||
self.assertEqual(tuple(masked_x.shape), (E, self.MAX_M, self.HIDDEN))
|
||||
|
||||
for e, (s, c) in enumerate(zip(starts, counts)):
|
||||
for j in range(c):
|
||||
torch.testing.assert_close(
|
||||
masked_x[e, j].float(), recv_x[s + j].float()
|
||||
)
|
||||
if with_scale:
|
||||
torch.testing.assert_close(
|
||||
masked_x_scale[e, j].float(), scale[s + j].float()
|
||||
)
|
||||
|
||||
weights = None
|
||||
if topk:
|
||||
weights = torch.zeros(total, dtype=torch.float32, device=DEVICE)
|
||||
for r in _real_rows(starts, counts):
|
||||
weights[r] = 0.25 + (r % 7) * 0.1
|
||||
out = masked_slab_to_expand(
|
||||
masked_x, psum, total, self.ALIGN, topk_weights=weights
|
||||
)
|
||||
self.assertEqual(tuple(out.shape), (total, self.HIDDEN))
|
||||
for e, (s, c) in enumerate(zip(starts, counts)):
|
||||
for j in range(c):
|
||||
expected = masked_x[e, j].float()
|
||||
if topk:
|
||||
expected = (expected * weights[s + j]).to(masked_x.dtype).float()
|
||||
torch.testing.assert_close(out[s + j].float(), expected)
|
||||
|
||||
def test_roundtrip_bf16(self):
|
||||
self._check_expand_roundtrip([3, 0, 5, 1], torch.bfloat16, with_scale=False)
|
||||
|
||||
def test_roundtrip_bf16_with_topk_weight(self):
|
||||
self._check_expand_roundtrip(
|
||||
[2, 4, 0, 7], torch.bfloat16, with_scale=False, topk=True
|
||||
)
|
||||
|
||||
def test_roundtrip_fp8_with_scale(self):
|
||||
self._check_expand_roundtrip([3, 1, 6, 2], torch.float8_e4m3fn, with_scale=True)
|
||||
|
||||
def test_empty_experts(self):
|
||||
self._check_expand_roundtrip([0, 0, 0, 0], torch.bfloat16, with_scale=False)
|
||||
|
||||
def test_single_hot_expert(self):
|
||||
self._check_expand_roundtrip(
|
||||
[0, self.MAX_M, 0, 0], torch.bfloat16, with_scale=False, topk=True
|
||||
)
|
||||
|
||||
def test_count_at_max_m_boundary(self):
|
||||
self._check_expand_roundtrip(
|
||||
[self.MAX_M, 1, self.MAX_M], torch.bfloat16, with_scale=False
|
||||
)
|
||||
|
||||
def test_overflow_fails_fast(self):
|
||||
counts = [self.MAX_M + 1, 2]
|
||||
recv_x, scale, psum, starts, total = _build_layout(
|
||||
counts, self.ALIGN, self.HIDDEN, torch.bfloat16
|
||||
)
|
||||
with self.assertRaises(RuntimeError):
|
||||
expand_to_masked_slab(
|
||||
recv_x, None, psum, len(counts), self.MAX_M, self.ALIGN
|
||||
)
|
||||
|
||||
def _production_packed_ue8m0_layout(self, counts):
|
||||
"""Build expanded rows with the production packed UE8M0 quantizer."""
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
sglang_per_token_group_quant_fp8,
|
||||
)
|
||||
|
||||
# hidden=1024 ensures the packed scale has multiple columns.
|
||||
hidden = 1024
|
||||
raw, _, psum, starts, total = _build_layout(
|
||||
counts, self.ALIGN, hidden, torch.bfloat16
|
||||
)
|
||||
recv_x, recv_x_scale = sglang_per_token_group_quant_fp8(
|
||||
raw,
|
||||
128,
|
||||
column_major_scales=True,
|
||||
scale_tma_aligned=True,
|
||||
scale_ue8m0=True,
|
||||
)
|
||||
self.assertEqual(recv_x_scale.dtype, torch.int32)
|
||||
self.assertGreater(recv_x_scale.shape[1], 1, "pack dim must be indexed")
|
||||
self.assertNotEqual(recv_x_scale.stride(1), 1)
|
||||
return recv_x, recv_x_scale, psum, starts, total, hidden
|
||||
|
||||
def test_fp8_packed_ue8m0_scale_from_production_quantizer(self):
|
||||
counts = [3, 1, 6, 2]
|
||||
recv_x, recv_x_scale, psum, starts, _, hidden = (
|
||||
self._production_packed_ue8m0_layout(counts)
|
||||
)
|
||||
E = len(counts)
|
||||
masked_x, masked_x_scale, masked_m = expand_to_masked_slab(
|
||||
recv_x, recv_x_scale, psum, E, self.MAX_M, self.ALIGN
|
||||
)
|
||||
self.assertEqual(masked_m.tolist(), list(counts))
|
||||
self.assertEqual(tuple(masked_x.shape), (E, self.MAX_M, hidden))
|
||||
for e, (s, c) in enumerate(zip(starts, counts)):
|
||||
for j in range(c):
|
||||
torch.testing.assert_close(
|
||||
masked_x[e, j].float(), recv_x[s + j].float()
|
||||
)
|
||||
torch.testing.assert_close(masked_x_scale[e, j], recv_x_scale[s + j])
|
||||
|
||||
def test_expand_under_cuda_graph_capture(self):
|
||||
# Exercise replay with the production packed scale layout.
|
||||
counts = [3, 1, 6, 2]
|
||||
recv_x, recv_x_scale, psum, starts, _, _ = self._production_packed_ue8m0_layout(
|
||||
counts
|
||||
)
|
||||
E = len(counts)
|
||||
warm = torch.cuda.Stream()
|
||||
warm.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(warm):
|
||||
expand_to_masked_slab(recv_x, recv_x_scale, psum, E, self.MAX_M, self.ALIGN)
|
||||
torch.cuda.current_stream().wait_stream(warm)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
masked_x, masked_x_scale, masked_m = expand_to_masked_slab(
|
||||
recv_x, recv_x_scale, psum, E, self.MAX_M, self.ALIGN
|
||||
)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(masked_m.tolist(), list(counts))
|
||||
for e, (s, c) in enumerate(zip(starts, counts)):
|
||||
for j in range(c):
|
||||
torch.testing.assert_close(
|
||||
masked_x[e, j].float(), recv_x[s + j].float()
|
||||
)
|
||||
torch.testing.assert_close(masked_x_scale[e, j], recv_x_scale[s + j])
|
||||
|
||||
|
||||
class TestDeepEPv2HandleLifecycle(CustomTestCase):
|
||||
"""CPU-only dispatch/combine handle guards."""
|
||||
|
||||
@staticmethod
|
||||
def _bare_impl():
|
||||
from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import _DeepEPv2Impl
|
||||
|
||||
impl = object.__new__(_DeepEPv2Impl)
|
||||
impl._handle = None
|
||||
impl._pad_empty_combine = False
|
||||
return impl
|
||||
|
||||
def test_combine_without_dispatch_raises(self):
|
||||
impl = self._bare_impl()
|
||||
with self.assertRaisesRegex(RuntimeError, "without a valid dispatch handle"):
|
||||
impl.combine(None)
|
||||
|
||||
def test_dispatch_with_unconsumed_handle_raises(self):
|
||||
impl = self._bare_impl()
|
||||
impl._handle = object()
|
||||
with self.assertRaisesRegex(RuntimeError, "unconsumed"):
|
||||
impl.dispatch(None, None)
|
||||
|
||||
def test_handle_cleared_when_combine_fails(self):
|
||||
impl = self._bare_impl()
|
||||
impl._handle = object()
|
||||
impl._pad_empty_combine = True
|
||||
|
||||
def _boom():
|
||||
raise RuntimeError("boom")
|
||||
|
||||
impl._get_buffer = _boom
|
||||
with self.assertRaisesRegex(RuntimeError, "boom"):
|
||||
impl.combine(None)
|
||||
self.assertIsNone(impl._handle)
|
||||
self.assertFalse(impl._pad_empty_combine)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,16 +1,13 @@
|
||||
"""The hpc_ops MoE runner backend makes the standard dispatcher keep global
|
||||
expert ids, so a quant method that silently falls back to another runner
|
||||
(e.g. an unquantized MoE) would misroute tokens under EP>1. MoeRunner must
|
||||
reject that combination loudly at startup.
|
||||
"""
|
||||
"""Startup guards for MoE runner and dispatcher quantization contracts."""
|
||||
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
|
||||
from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
||||
from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend
|
||||
from sglang.srt.runtime_context import get_flags
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
@@ -18,27 +15,27 @@ register_cpu_ci(est_time=6, suite="base-c-test-cpu")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _runner_backend_flag():
|
||||
def _moe_flags():
|
||||
moe = get_flags().moe
|
||||
saved = moe.runner_backend
|
||||
saved = (moe.runner_backend, moe.a2a_backend)
|
||||
yield moe
|
||||
moe.runner_backend = saved
|
||||
moe.runner_backend, moe.a2a_backend = saved
|
||||
|
||||
|
||||
def test_non_hpc_runner_rejected_when_hpc_ops_requested(_runner_backend_flag):
|
||||
_runner_backend_flag.runner_backend = MoeRunnerBackend.HPC_OPS
|
||||
def test_non_hpc_runner_rejected_when_hpc_ops_requested(_moe_flags):
|
||||
_moe_flags.runner_backend = MoeRunnerBackend.HPC_OPS
|
||||
with pytest.raises(ValueError, match="hpc_ops"):
|
||||
MoeRunner(MoeRunnerBackend.TRITON, MoeRunnerConfig())
|
||||
|
||||
|
||||
def test_triton_runner_allowed_without_hpc_ops(_runner_backend_flag):
|
||||
_runner_backend_flag.runner_backend = MoeRunnerBackend.TRITON
|
||||
def test_triton_runner_allowed_without_hpc_ops(_moe_flags):
|
||||
_moe_flags.runner_backend = MoeRunnerBackend.TRITON
|
||||
runner = MoeRunner(MoeRunnerBackend.TRITON, MoeRunnerConfig())
|
||||
assert runner.runner_core is not None
|
||||
|
||||
|
||||
def test_direct_kernel_quant_method_rejected_when_hpc_ops_requested(
|
||||
_runner_backend_flag,
|
||||
_moe_flags,
|
||||
):
|
||||
# W4AFp8MoEMethod never constructs a MoeRunner (apply() calls its kernel
|
||||
# directly), so it bypasses the MoeRunner-level guard; the layer-level
|
||||
@@ -49,15 +46,95 @@ def test_direct_kernel_quant_method_rejected_when_hpc_ops_requested(
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod
|
||||
from sglang.srt.layers.quantization.w4afp8 import W4AFp8MoEMethod
|
||||
|
||||
_runner_backend_flag.runner_backend = MoeRunnerBackend.HPC_OPS
|
||||
_moe_flags.runner_backend = MoeRunnerBackend.HPC_OPS
|
||||
with pytest.raises(ValueError, match="hpc_ops"):
|
||||
_validate_hpc_ops_quant_method(object.__new__(W4AFp8MoEMethod))
|
||||
# The FP8 method (the one the hpc_ops runner supports) passes.
|
||||
_validate_hpc_ops_quant_method(object.__new__(Fp8MoEMethod))
|
||||
# Without hpc_ops requested, any quant method passes.
|
||||
_runner_backend_flag.runner_backend = MoeRunnerBackend.TRITON
|
||||
_moe_flags.runner_backend = MoeRunnerBackend.TRITON
|
||||
_validate_hpc_ops_quant_method(object.__new__(W4AFp8MoEMethod))
|
||||
|
||||
|
||||
def _fp8_method(**overrides):
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8MoEMethod
|
||||
|
||||
values = {
|
||||
"activation_scheme": "dynamic",
|
||||
"weight_block_size": (128, 128),
|
||||
"use_mxfp8": False,
|
||||
"is_fp4_expert": False,
|
||||
}
|
||||
values.update(overrides)
|
||||
method = object.__new__(Fp8MoEMethod)
|
||||
method.quant_config = SimpleNamespace(
|
||||
activation_scheme=values["activation_scheme"],
|
||||
)
|
||||
method.weight_block_size = values["weight_block_size"]
|
||||
method.use_mxfp8 = values["use_mxfp8"]
|
||||
method.is_fp4_expert = values["is_fp4_expert"]
|
||||
return method
|
||||
|
||||
|
||||
def test_deepep_v2_quant_contract_accepts_blockwise_fp8(_moe_flags):
|
||||
from sglang.srt.layers.moe.fused_moe_triton.layer import (
|
||||
_validate_deepep_v2_quant_method,
|
||||
)
|
||||
|
||||
_moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2
|
||||
_validate_deepep_v2_quant_method(_fp8_method(weight_block_size=[128, 128]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("overrides", "expected"),
|
||||
[
|
||||
({"activation_scheme": "static"}, "activation_scheme"),
|
||||
({"weight_block_size": None}, "weight_block_size"),
|
||||
({"weight_block_size": (1, 32), "use_mxfp8": True}, "MXFP8"),
|
||||
({"is_fp4_expert": True}, "FP4 experts"),
|
||||
],
|
||||
)
|
||||
def test_deepep_v2_quant_contract_rejects_incompatible_fp8(
|
||||
_moe_flags, overrides, expected
|
||||
):
|
||||
from sglang.srt.layers.moe.fused_moe_triton.layer import (
|
||||
_validate_deepep_v2_quant_method,
|
||||
)
|
||||
|
||||
_moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
_validate_deepep_v2_quant_method(_fp8_method(**overrides))
|
||||
|
||||
|
||||
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))
|
||||
|
||||
|
||||
def test_deepep_v2_quant_contract_does_not_affect_other_backends(_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
|
||||
|
||||
_moe_flags.a2a_backend = MoeA2ABackend.DEEPEP
|
||||
_validate_deepep_v2_quant_method(object.__new__(UnquantizedFusedMoEMethod))
|
||||
|
||||
|
||||
def test_deepep_v2_runner_backstop(_moe_flags):
|
||||
_moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2
|
||||
with pytest.raises(ValueError, match="deep_gemm"):
|
||||
MoeRunner(MoeRunnerBackend.TRITON, MoeRunnerConfig())
|
||||
assert MoeRunner(MoeRunnerBackend.DEEP_GEMM, MoeRunnerConfig()).runner_core
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -2106,6 +2106,328 @@ class TestSamplingBackendTokenOracleEnvGate(CustomTestCase):
|
||||
self.assertEqual(parsed.sampling_backend, "token_oracle")
|
||||
|
||||
|
||||
class TestDeepEPv2Args(CustomTestCase):
|
||||
"""DeepEP v2 server-argument resolution and validation."""
|
||||
|
||||
def _args(self, **overrides):
|
||||
server_args = ServerArgs(model_path="dummy", moe_a2a_backend="deepep_v2")
|
||||
server_args.model_config = SimpleNamespace(
|
||||
hf_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"])
|
||||
)
|
||||
# The dummy path does not initialize phase configs.
|
||||
server_args.cuda_graph_config = CudaGraphConfig(
|
||||
decode=PhaseConfig(backend=Backend.FULL, max_bs=512),
|
||||
prefill=PhaseConfig(backend=Backend.FULL, max_bs=512),
|
||||
)
|
||||
server_args._resolved_overrides = []
|
||||
valid = {f.name for f in dataclasses.fields(ServerArgs)}
|
||||
for key, value in overrides.items():
|
||||
# Reject stale field names before setattr silently accepts them.
|
||||
assert key in valid, f"{key} is not a ServerArgs field"
|
||||
setattr(server_args, key, value)
|
||||
return server_args
|
||||
|
||||
def test_validated_architectures_allowed(self):
|
||||
for architecture in (
|
||||
"DeepseekV3ForCausalLM",
|
||||
"DeepseekV4ForCausalLM",
|
||||
"Qwen3MoeForCausalLM",
|
||||
):
|
||||
args = self._args(moe_runner_backend="deep_gemm")
|
||||
args.model_config.hf_config.architectures = [architecture]
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_unvalidated_and_missing_architectures_rejected(self):
|
||||
for architectures in (
|
||||
["Qwen2MoeForCausalLM"],
|
||||
["Qwen3_5MoeForCausalLM"],
|
||||
[],
|
||||
None,
|
||||
):
|
||||
args = self._args(moe_runner_backend="deep_gemm")
|
||||
args.model_config.hf_config.architectures = architectures
|
||||
with self.assertRaisesRegex(ValueError, "not validated"):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_instance_connector_rejected(self):
|
||||
args = self._args(
|
||||
model_path="instance://worker/model",
|
||||
moe_runner_backend="deep_gemm",
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "instance connector"):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_deterministic_inference_rejected(self):
|
||||
args = self._args(
|
||||
moe_runner_backend="deep_gemm",
|
||||
enable_deterministic_inference=True,
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "deterministic sorting"):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_rl_on_policy_deterministic_inference_rejected(self):
|
||||
args = self._args(
|
||||
moe_runner_backend="deep_gemm",
|
||||
rl_on_policy_target="fsdp",
|
||||
)
|
||||
args.model_config.hf_config.architectures = ["Qwen3MoeForCausalLM"]
|
||||
with (
|
||||
envs.SGLANG_VLM_CACHE_SIZE_MB.override(envs.SGLANG_VLM_CACHE_SIZE_MB.get()),
|
||||
envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.override(
|
||||
envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()
|
||||
),
|
||||
):
|
||||
args._handle_deterministic_inference()
|
||||
with self.assertRaisesRegex(ValueError, "deterministic sorting"):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_deterministic_inference_does_not_affect_legacy_deepep(self):
|
||||
args = self._args(
|
||||
moe_a2a_backend="deepep",
|
||||
moe_runner_backend="deep_gemm",
|
||||
enable_deterministic_inference=True,
|
||||
)
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_runner_restored_by_declaration_fails_fast(self):
|
||||
# Validate the declaration-resolved runner rather than the raw field.
|
||||
args = self._args(moe_runner_backend="auto")
|
||||
args._resolved_overrides = [
|
||||
("test_mxfp8", {"moe_runner_backend": "flashinfer_trtllm"})
|
||||
]
|
||||
with self.assertRaises(ValueError):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_declarations_resolve_ep_size_and_fusion(self):
|
||||
from sglang.srt.arg_groups.overrides import resolved_view
|
||||
|
||||
args = self._args(moe_runner_backend="auto", tp_size=2)
|
||||
args._handle_a2a_moe()
|
||||
self.assertEqual(resolved_view(args).ep_size, args.tp_size)
|
||||
self.assertTrue(resolved_view(args).disable_shared_experts_fusion)
|
||||
|
||||
def test_auto_runner_defaults_to_deep_gemm(self):
|
||||
from sglang.srt.arg_groups.overrides import resolved_view
|
||||
|
||||
args = self._args(moe_runner_backend="auto")
|
||||
args._handle_a2a_moe()
|
||||
self.assertEqual(resolved_view(args).moe_runner_backend, "deep_gemm")
|
||||
|
||||
def test_unsupported_runner_rejected(self):
|
||||
args = self._args(moe_runner_backend="flashinfer_trtllm")
|
||||
with self.assertRaises(ValueError):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_triton_runner_rejected(self):
|
||||
args = self._args(moe_runner_backend="triton")
|
||||
with self.assertRaises(ValueError):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_decode_graph_stays_enabled_in_both_comm_modes(self):
|
||||
for mode in ("direct", "hybrid"):
|
||||
args = self._args(moe_runner_backend="deep_gemm", deepep_v2_mode=mode)
|
||||
args._handle_a2a_moe()
|
||||
self.assertEqual(args.cuda_graph_config.decode.backend, Backend.FULL)
|
||||
self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED)
|
||||
|
||||
def test_two_batch_overlap_rejected(self):
|
||||
args = self._args(moe_runner_backend="deep_gemm", enable_two_batch_overlap=True)
|
||||
with self.assertRaises(ValueError):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_speculative_draft_backend_rejected(self):
|
||||
for main_backend in ("none", "deepep", "deepep_v2"):
|
||||
args = self._args(
|
||||
moe_a2a_backend=main_backend,
|
||||
moe_runner_backend="deep_gemm",
|
||||
speculative_moe_a2a_backend="deepep_v2",
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "speculative draft backend"):
|
||||
args._validate_deepep_v2_speculative_draft()
|
||||
|
||||
def test_inherited_speculative_draft_backend_rejected(self):
|
||||
args = self._args(
|
||||
moe_runner_backend="deep_gemm",
|
||||
speculative_algorithm="EAGLE",
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "speculative draft backend"):
|
||||
args._validate_deepep_v2_speculative_draft()
|
||||
|
||||
def test_ngram_does_not_inherit_a_draft_backend(self):
|
||||
args = self._args(
|
||||
moe_runner_backend="deep_gemm",
|
||||
speculative_algorithm="NGRAM",
|
||||
)
|
||||
args._validate_deepep_v2_speculative_draft()
|
||||
|
||||
def test_explicit_legacy_speculative_backend_allowed(self):
|
||||
args = self._args(
|
||||
moe_runner_backend="deep_gemm",
|
||||
speculative_algorithm="EAGLE",
|
||||
speculative_moe_a2a_backend="deepep",
|
||||
)
|
||||
args._validate_deepep_v2_speculative_draft()
|
||||
|
||||
def test_resolved_legacy_speculative_backend_allowed(self):
|
||||
args = self._args(
|
||||
moe_runner_backend="deep_gemm",
|
||||
speculative_algorithm="EAGLE",
|
||||
)
|
||||
args._resolved_overrides = [
|
||||
(
|
||||
"test_speculative_backend",
|
||||
{"speculative_moe_a2a_backend": "deepep"},
|
||||
)
|
||||
]
|
||||
args._validate_deepep_v2_speculative_draft()
|
||||
|
||||
def test_prefill_chunk_exceeding_cap_rejected(self):
|
||||
args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=2048)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1024):
|
||||
with self.assertRaisesRegex(ValueError, "NUM_MAX_DISPATCH_TOKENS_PER_RANK"):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_prefill_chunk_at_cap_boundary_accepted(self):
|
||||
args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=1024)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1024):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_dynamic_chunking_probe_is_included(self):
|
||||
args = self._args(
|
||||
chunked_prefill_size=1024,
|
||||
max_prefill_tokens=1024,
|
||||
enable_dynamic_chunking=True,
|
||||
pp_size=2,
|
||||
disaggregation_mode="prefill",
|
||||
)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1024):
|
||||
with self.assertRaisesRegex(ValueError, "required=1280"):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_disabled_chunking_uses_max_prefill_tokens(self):
|
||||
for disabled in (None, 0, -1):
|
||||
args = self._args(
|
||||
chunked_prefill_size=disabled,
|
||||
max_prefill_tokens=1024,
|
||||
disaggregation_mode="prefill",
|
||||
)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128):
|
||||
with self.assertRaisesRegex(ValueError, "required=1024"):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_decode_role_skips_prefill_capacity(self):
|
||||
args = self._args(
|
||||
chunked_prefill_size=4096,
|
||||
disaggregation_mode="decode",
|
||||
max_running_requests=32,
|
||||
dp_size=1,
|
||||
)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_decode_graph_capacity_boundaries(self):
|
||||
for max_bs, raises in ((128, False), (129, True)):
|
||||
args = self._args(
|
||||
disaggregation_mode="decode",
|
||||
max_running_requests=None,
|
||||
)
|
||||
args.cuda_graph_config.decode.max_bs = max_bs
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128):
|
||||
if raises:
|
||||
with self.assertRaisesRegex(ValueError, "decode CUDA graph"):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
else:
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_dp_attention_divides_max_running_requests_per_rank(self):
|
||||
args = self._args(
|
||||
disaggregation_mode="decode",
|
||||
max_running_requests=256,
|
||||
tp_size=8,
|
||||
dp_size=8,
|
||||
enable_dp_attention=True,
|
||||
)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_tp_only_max_running_requests_is_not_divided(self):
|
||||
args = self._args(
|
||||
disaggregation_mode="decode",
|
||||
max_running_requests=256,
|
||||
tp_size=8,
|
||||
dp_size=1,
|
||||
enable_dp_attention=False,
|
||||
)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128):
|
||||
with self.assertRaisesRegex(ValueError, "decode CUDA graph"):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_memory_derived_eager_pool_remains_runtime_validated(self):
|
||||
args = self._args(
|
||||
disaggregation_mode="decode",
|
||||
max_running_requests=None,
|
||||
)
|
||||
args.cuda_graph_config.decode.backend = Backend.DISABLED
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_speculative_decode_width_is_included(self):
|
||||
args = self._args(
|
||||
disaggregation_mode="decode",
|
||||
speculative_algorithm="EAGLE",
|
||||
speculative_num_draft_tokens=8,
|
||||
max_running_requests=256,
|
||||
dp_size=8,
|
||||
enable_dp_attention=True,
|
||||
)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128):
|
||||
with self.assertRaisesRegex(ValueError, "tokens/request=8"):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_adaptive_speculative_uses_widest_candidate(self):
|
||||
args = self._args(
|
||||
disaggregation_mode="decode",
|
||||
speculative_algorithm="EAGLE",
|
||||
speculative_num_draft_tokens=4,
|
||||
speculative_adaptive=True,
|
||||
max_running_requests=128,
|
||||
dp_size=8,
|
||||
enable_dp_attention=True,
|
||||
)
|
||||
with patch.object(
|
||||
ServerArgs,
|
||||
"max_speculative_num_draft_tokens",
|
||||
new=property(lambda _self: 16),
|
||||
):
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128):
|
||||
with self.assertRaisesRegex(ValueError, "tokens/request=16"):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_prefill_role_skips_decode_capacity(self):
|
||||
args = self._args(
|
||||
disaggregation_mode="prefill",
|
||||
chunked_prefill_size=64,
|
||||
max_running_requests=8192,
|
||||
)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(128):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_other_backend_skips_capacity_validation(self):
|
||||
args = self._args(
|
||||
moe_a2a_backend="deepep",
|
||||
chunked_prefill_size=4096,
|
||||
max_running_requests=4096,
|
||||
)
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
def test_capacity_validation_uses_resolved_backend(self):
|
||||
args = self._args(chunked_prefill_size=4096)
|
||||
args._resolved_overrides = [("test", {"moe_a2a_backend": "deepep"})]
|
||||
with envs.SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK.override(1):
|
||||
args._validate_deepep_v2_dispatch_token_budget()
|
||||
|
||||
|
||||
class TestHandleCrashDumpEnv(CustomTestCase):
|
||||
_COREDUMP_ENV_KEYS = (
|
||||
"CUDA_ENABLE_COREDUMP_ON_EXCEPTION",
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
"""DeepEP-family backend recognition in RoutedExpertsCapturer."""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.utils import MoeA2ABackend
|
||||
from sglang.srt.state_capturer import routed_experts as re_mod
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestScatteredA2ABackendHelper(CustomTestCase):
|
||||
def test_classification(self):
|
||||
expected = {
|
||||
"deepep": True,
|
||||
"deepep_v2": True,
|
||||
"none": False,
|
||||
"mooncake": False,
|
||||
}
|
||||
for value, exp in expected.items():
|
||||
with mock.patch.object(
|
||||
re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend(value)
|
||||
):
|
||||
self.assertEqual(
|
||||
re_mod._is_scattered_a2a_backend(), exp, f"backend={value}"
|
||||
)
|
||||
|
||||
|
||||
class TestGetLocalSliceBackendBranch(CustomTestCase):
|
||||
T, L, K = 16, 3, 4
|
||||
|
||||
def _capturer(self):
|
||||
cap = object.__new__(re_mod.RoutedExpertsCapturer)
|
||||
buf = torch.arange(self.T * self.L * self.K, dtype=torch.int32).reshape(
|
||||
self.T, self.L, self.K
|
||||
)
|
||||
cap.device_cache = SimpleNamespace(buffer=buf)
|
||||
cap.topk_size = self.K
|
||||
return cap, buf
|
||||
|
||||
def _slice(self, cap, n_local):
|
||||
fb = SimpleNamespace(out_cache_loc=torch.empty(n_local))
|
||||
return cap._get_local_slice(fb, can_run_graph=False, cuda_graph_batch=None)
|
||||
|
||||
def test_deepep_v2_reads_buffer_head(self):
|
||||
cap, buf = self._capturer()
|
||||
with mock.patch.object(
|
||||
re_mod, "is_dp_attention_enabled", return_value=True
|
||||
), mock.patch.object(
|
||||
re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("deepep_v2")
|
||||
):
|
||||
out = self._slice(cap, n_local=5)
|
||||
self.assertTrue(torch.equal(out, buf[0:5, :, : self.K]))
|
||||
|
||||
def test_deepep_v2_matches_deepep(self):
|
||||
cap, _ = self._capturer()
|
||||
outs = []
|
||||
for backend in ("deepep", "deepep_v2"):
|
||||
with mock.patch.object(
|
||||
re_mod, "is_dp_attention_enabled", return_value=True
|
||||
), mock.patch.object(
|
||||
re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend(backend)
|
||||
):
|
||||
outs.append(self._slice(cap, n_local=7))
|
||||
self.assertTrue(torch.equal(outs[0], outs[1]))
|
||||
|
||||
def test_tp_moe_reads_global_offset(self):
|
||||
cap, buf = self._capturer()
|
||||
with mock.patch.object(
|
||||
re_mod, "is_dp_attention_enabled", return_value=True
|
||||
), mock.patch.object(
|
||||
re_mod, "get_moe_a2a_backend", return_value=MoeA2ABackend("none")
|
||||
), mock.patch.object(
|
||||
re_mod, "get_dp_local_slice_cpu", return_value=(6, 4)
|
||||
):
|
||||
out = self._slice(cap, n_local=999)
|
||||
self.assertTrue(torch.equal(out, buf[6:10, :, : self.K]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user