[Feature] Add DeepEPv2 (ElasticBuffer) MoE A2A backend (#29525)
Co-authored-by: menyu <menyu@nvidia.com>
This commit is contained in:
MengYu
@@ -0,0 +1,207 @@
|
||||
"""DP>1 readback of routed experts over DeepEP-class a2a backends.
|
||||
|
||||
With DP attention + a DeepEP-class a2a backend, the MoE layer sees only the
|
||||
attention rank's DP-local tokens, so RoutedExpertsCapturer must gather at
|
||||
capture time and read back from the buffer head. If the backend is not
|
||||
recognized, requests owned by dp_rank > 0 read unwritten buffer rows and
|
||||
silently return garbage expert ids (dp_rank 0 sits at offset 0 and looks
|
||||
correct, which is why a DP>1 test is required).
|
||||
|
||||
Oracle: solo-vs-concurrent consistency. A request served alone is correct
|
||||
even on a misclassifying tree (with the other rank empty, the global offset
|
||||
degenerates to 0), so its per-token expert sets form a valid baseline. The
|
||||
same prompts served concurrently must reproduce those sets; a misclassified
|
||||
backend instead reads whatever the offset region holds (often well-formed
|
||||
rows belonging to other tokens or graph warmup, which per-row validity
|
||||
checks cannot catch). Radix cache is disabled so the concurrent phase cannot
|
||||
serve cached prefix rows written by the solo phase.
|
||||
|
||||
Uses a dummy-weight single-layer 24-expert DeepSeek-V3 so each server boots
|
||||
in seconds (same pattern as test_deepseek_v3_cutedsl_4gpu.py); generation
|
||||
quality is irrelevant — only the capture/readback plumbing is under test.
|
||||
"""
|
||||
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import pybase64
|
||||
import requests
|
||||
|
||||
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="deepep-8-gpu-h200")
|
||||
|
||||
_MODEL = os.environ.get("SGLANG_ROUTED_EXPERTS_TEST_MODEL", "deepseek-ai/DeepSeek-V3")
|
||||
_NUM_EXPERTS = 24
|
||||
_NUM_LAYERS = 1
|
||||
_TOPK = 8 # DeepSeek-V3 num_experts_per_tok
|
||||
|
||||
_DUMMY_WEIGHT_ENV = {
|
||||
# Dummy random weights legitimately produce NaN logits; sanitize instead
|
||||
# of crashing (same rationale as test_deepseek_v3_cutedsl_4gpu.py).
|
||||
"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)
|
||||
|
||||
|
||||
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",
|
||||
"--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):
|
||||
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):
|
||||
# Phase 1 — solo baselines: sequential requests leave the other DP
|
||||
# rank empty, the global offset degenerates to 0, and the readback is
|
||||
# correct even when the backend is misclassified.
|
||||
solo = [self._one_request(i) for i in range(self._N_REQ)]
|
||||
|
||||
# Phase 2 — the same prompts concurrently: joint forward batches give
|
||||
# dp_rank > 0 requests a non-zero global offset, which is exactly the
|
||||
# path a misclassified backend gets wrong.
|
||||
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"
|
||||
)
|
||||
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,277 @@
|
||||
"""Unit tests for the DeepEP v2 masked-masked_x repack Triton kernels.
|
||||
|
||||
Covers the corner cases flagged in review: empty expert, single hot expert,
|
||||
per-expert count near / over max_m (overflow -> fail-fast, not silent truncation),
|
||||
top-k weight fusion on real rows only, expanded<->masked_x round-trip layout, and the
|
||||
fp8 activation+scale path.
|
||||
"""
|
||||
|
||||
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 (recv_x, recv_x_scale, psum, starts, total) for given per-expert counts.
|
||||
|
||||
Mirrors the DeepEP v2 expanded layout: expert e occupies rows
|
||||
[align(psum[e-1]), psum[e]) with psum[-1] == 0.
|
||||
"""
|
||||
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)
|
||||
|
||||
# Unique value per real row (kept small so the x2 column stays inside the
|
||||
# e4m3 range), alternating x1/x2 along hidden so a kernel that broadcast one
|
||||
# column across the row, or mis-strided the hidden offset, is caught.
|
||||
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)
|
||||
# Distinct value per scale column: a row constant along scale_hidden
|
||||
# cannot catch a kernel that reads column 0 for every column.
|
||||
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
|
||||
)
|
||||
|
||||
# masked_m == real per-expert count
|
||||
self.assertEqual(masked_m.tolist(), list(counts))
|
||||
self.assertEqual(tuple(masked_x.shape), (E, self.MAX_M, self.HIDDEN))
|
||||
|
||||
# masked_x real rows == source expanded rows
|
||||
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()
|
||||
)
|
||||
|
||||
# round-trip back to expanded order
|
||||
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):
|
||||
# exactly max_m must be kept (no overflow, no truncation)
|
||||
self._check_expand_roundtrip(
|
||||
[self.MAX_M, 1, self.MAX_M], torch.bfloat16, with_scale=False
|
||||
)
|
||||
|
||||
def test_overflow_fails_fast(self):
|
||||
# one expert exceeds max_m -> must raise, not silently truncate
|
||||
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):
|
||||
"""Expanded FP8 rows + scales built by the PRODUCTION quantizer with the
|
||||
Blackwell flags (packed ue8m0, column-major): scale is int32 with
|
||||
pack-dim stride != 1, unlike Hopper's row-major fp32."""
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
sglang_per_token_group_quant_fp8,
|
||||
)
|
||||
|
||||
# 1024 = 8 quant groups of 128 -> the packed scale has ceil(8/4) = 2 int32
|
||||
# columns. At hidden <= 512 it collapses to a single column: the pack-dim
|
||||
# offset is always 0 and the pack-dim stride is never exercised.
|
||||
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):
|
||||
# The Blackwell dispatch scale is packed ue8m0 in column-major layout,
|
||||
# so the repack must honor the scale pack-dim stride; the row-major
|
||||
# fp32 cases above (stride(1) == 1) cannot regress it.
|
||||
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):
|
||||
# The masked repack runs inside the captured decode CUDA graph, so it
|
||||
# must be capture-safe (no host sync) and correct after replay, with the
|
||||
# production packed ue8m0 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 guards of the dispatch/combine handle lifecycle.
|
||||
|
||||
The guards are ordered before any DeepEP work, so misuse is testable
|
||||
without deep_ep installed and without a GPU. The positive dispatch ->
|
||||
combine path needs real ElasticBuffer communication and is covered by the
|
||||
GPU accuracy runs instead.
|
||||
"""
|
||||
|
||||
@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()
|
||||
@@ -1898,6 +1898,120 @@ class TestSamplingBackendTokenOracleEnvGate(CustomTestCase):
|
||||
self.assertEqual(parsed.sampling_backend, "token_oracle")
|
||||
|
||||
|
||||
class TestDeepEPv2Args(CustomTestCase):
|
||||
"""DeepEP v2 server-args resolution + validation. The dummy-model path
|
||||
short-circuits __post_init__, so _handle_a2a_moe() is invoked directly."""
|
||||
|
||||
def _args(self, **overrides):
|
||||
server_args = ServerArgs(model_path="dummy", moe_a2a_backend="deepep_v2")
|
||||
# The deepep_v2 branch mutates cuda_graph_config.{decode,prefill}.backend,
|
||||
# so it must exist (the dummy path leaves it unset otherwise).
|
||||
server_args.cuda_graph_config = CudaGraphConfig(
|
||||
decode=PhaseConfig(backend=Backend.FULL, max_bs=512),
|
||||
prefill=PhaseConfig(backend=Backend.FULL, max_bs=512),
|
||||
)
|
||||
valid = {f.name for f in dataclasses.fields(ServerArgs)}
|
||||
for key, value in overrides.items():
|
||||
# ServerArgs has no __slots__, so setattr of a stale field name would
|
||||
# silently succeed and leave the test asserting nothing.
|
||||
assert key in valid, f"{key} is not a ServerArgs field"
|
||||
setattr(server_args, key, value)
|
||||
return server_args
|
||||
|
||||
def test_runner_restored_by_declaration_fails_fast(self):
|
||||
# mxfp8 + auto: a model declaration restores an unsupported runner at
|
||||
# materialize time, which runs after this handler. The handler must
|
||||
# validate the declaration-resolved runner, not the raw value it just set.
|
||||
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_materialize_ep_size_and_fusion(self):
|
||||
from sglang.srt.arg_groups.overrides import materialize_declarations
|
||||
|
||||
args = self._args(moe_runner_backend="auto", tp_size=2)
|
||||
args._handle_a2a_moe()
|
||||
# ep_size / shared-experts fusion are declared by the a2a passes and land
|
||||
# on the fields only at materialization, like every other a2a backend.
|
||||
materialize_declarations(args)
|
||||
self.assertEqual(args.ep_size, args.tp_size)
|
||||
self.assertTrue(args.disable_shared_experts_fusion)
|
||||
|
||||
def test_auto_runner_defaults_to_deep_gemm(self):
|
||||
args = self._args(moe_runner_backend="auto")
|
||||
args._handle_a2a_moe()
|
||||
self.assertEqual(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):
|
||||
# deepep_v2 registers permute adapters for deep_gemm only. Rejecting
|
||||
# triton here is what keeps a user from reaching the permute registry
|
||||
# and dying on a bare assert inside the MoE forward.
|
||||
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):
|
||||
# Capturability follows the inference phase (masked decode), not the
|
||||
# comm mode, so neither direct nor hybrid may disable the decode graph.
|
||||
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()
|
||||
|
||||
# --- prefill capacity pre-check (per-rank chunk vs dispatch buffer cap) ---
|
||||
_CAP_ENV = "SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK"
|
||||
|
||||
def test_prefill_chunk_exceeding_cap_rejected(self):
|
||||
args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=2048)
|
||||
with patch.dict(os.environ, {self._CAP_ENV: "1024"}):
|
||||
with self.assertRaisesRegex(ValueError, "NUM_MAX_DISPATCH_TOKENS_PER_RANK"):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_prefill_chunk_at_cap_boundary_accepted(self):
|
||||
# chunk == cap is the documented (and currently benchmarked) edge; the
|
||||
# guard must be strict-greater-than.
|
||||
args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=1024)
|
||||
with patch.dict(os.environ, {self._CAP_ENV: "1024"}):
|
||||
args._handle_a2a_moe()
|
||||
self.assertEqual(args.moe_runner_backend, "deep_gemm")
|
||||
|
||||
def test_prefill_chunk_rejected_under_default_cap(self):
|
||||
# Default cap is 128: a typical 1024-token per-rank chunk must be
|
||||
# rejected at boot instead of at the first full prefill chunk.
|
||||
args = self._args(moe_runner_backend="deep_gemm", chunked_prefill_size=1024)
|
||||
with self.assertRaisesRegex(ValueError, "chunked prefill budget"):
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_prefill_chunk_check_skipped_for_decode_disaggregation(self):
|
||||
args = self._args(
|
||||
moe_runner_backend="deep_gemm",
|
||||
chunked_prefill_size=4096,
|
||||
disaggregation_mode="decode",
|
||||
)
|
||||
args._handle_a2a_moe()
|
||||
|
||||
def test_prefill_chunk_check_skipped_when_chunking_disabled(self):
|
||||
for disabled in (None, 0, -1):
|
||||
args = self._args(
|
||||
moe_runner_backend="deep_gemm", chunked_prefill_size=disabled
|
||||
)
|
||||
args._handle_a2a_moe()
|
||||
|
||||
|
||||
class TestHandleCrashDumpEnv(CustomTestCase):
|
||||
_COREDUMP_ENV_KEYS = (
|
||||
"CUDA_ENABLE_COREDUMP_ON_EXCEPTION",
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
"""DeepEP-class backend recognition in RoutedExpertsCapturer.
|
||||
|
||||
The capturer keys its buffer layout on the a2a backend: DeepEP-class
|
||||
dispatchers hand the MoE layer only the attention rank's DP-local tokens, so
|
||||
``capture()`` must attn-TP-gather and ``_get_local_slice()`` must read the
|
||||
buffer head instead of the global DP offset. These tests pin that DeepEP v2
|
||||
is classified like DeepEP (it shares that token topology); a miss makes
|
||||
dp_rank > 0 read unwritten rows (silent wrong data), see the DP>1 readback
|
||||
test in test/registered/ep/test_routed_experts_dp_readback.py.
|
||||
"""
|
||||
|
||||
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_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
class TestScatteredA2ABackendHelper(CustomTestCase):
|
||||
def test_classification(self):
|
||||
# deepep_v2 shares DeepEP's scattered token topology. Other backends
|
||||
# keep their existing classification (mooncake/mori are deliberately
|
||||
# not reclassified here).
|
||||
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 # buffer tokens, layers, top-k
|
||||
|
||||
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