minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Revert "[Feature] Add DeepEPv2 (ElasticBuffer) MoE A2A backend" (#35568)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-19 14:14:38 -07:00
committed by GitHub
parent a6bc0532c9
commit 1270204d2c
18 changed files with 21 additions and 2112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/ep/test_routed_experts_dp_readback.py
-207
View File
@@ -1,207 +0,0 @@
"""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()
test/registered/unit/layers/moe/test_deepep_v2_masked_slab.py
-277
View File
@@ -1,277 +0,0 @@
"""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()
test/registered/unit/server_args/test_server_args.py
-114
View File
@@ -1917,120 +1917,6 @@ 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",
test/registered/unit/state_capturer/test_routed_experts_scattered_a2a.py
-98
View File
@@ -1,98 +0,0 @@
"""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()