[RL] DeepEP support for --enable-return-routed-experts (#16859)
Co-authored-by: hnyls2002 <lsyincs@gmail.com> Co-authored-by: Junrong Lin <33685709+ocss884@users.noreply.github.com>
This commit is contained in:
Night
co-authored by
hnyls2002
Junrong Lin
hnyls2002
Junrong Lin
parent
d7385b575f
commit
b2420d72ff
@@ -6,10 +6,13 @@ import torch
|
||||
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.layers.dp_attention import (
|
||||
attn_tp_all_gather_into_tensor,
|
||||
get_attention_dp_rank,
|
||||
get_attention_tp_size,
|
||||
get_dp_local_info,
|
||||
is_dp_attention_enabled,
|
||||
)
|
||||
from sglang.srt.layers.moe import get_moe_a2a_backend
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.state_capturer.base import BaseTopkCapturer
|
||||
@@ -75,13 +78,40 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
|
||||
device_topk_size=topk_size + num_fused_shared_experts,
|
||||
)
|
||||
|
||||
# DeepEP a2a path: each attn-TP rank only sees its scattered slice of
|
||||
# topk_ids. All-gather across attn-TP at capture time so device_cache
|
||||
# holds the full batch and the existing _get_local_slice / D2H sync
|
||||
# paths work unchanged. Pre-allocate the gather target.
|
||||
if get_moe_a2a_backend().is_deepep():
|
||||
attn_tp_size = get_attention_tp_size() if is_dp_attention_enabled() else 1
|
||||
self.gather_buffer = torch.empty(
|
||||
(
|
||||
self.device_cache.buffer.shape[0] * attn_tp_size,
|
||||
self.device_cache.buffer.shape[2],
|
||||
),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def capture(self, layer_id: int, topk_indices: torch.Tensor):
|
||||
if get_moe_a2a_backend().is_deepep():
|
||||
local_topk = topk_indices
|
||||
topk_indices = self.gather_buffer[
|
||||
: local_topk.size(0) * get_attention_tp_size()
|
||||
]
|
||||
attn_tp_all_gather_into_tensor(topk_indices, local_topk)
|
||||
super().capture(layer_id, topk_indices)
|
||||
|
||||
def _get_local_slice(
|
||||
self,
|
||||
forward_batch: ForwardBatch,
|
||||
can_run_graph: bool,
|
||||
cuda_graph_batch: Optional[int],
|
||||
) -> torch.Tensor:
|
||||
if is_dp_attention_enabled():
|
||||
# Under DeepEP, capture() already attn_tp_all_gathered into the head of
|
||||
# the per-rank buffer, so the local DP rank's data lives at [0:N_local]
|
||||
# rather than at the global [start_pos:end_pos] offset.
|
||||
if is_dp_attention_enabled() and not get_moe_a2a_backend().is_deepep():
|
||||
local_start_pos, local_num_tokens = get_dp_local_info(forward_batch)
|
||||
if can_run_graph:
|
||||
local_start_pos = get_attention_dp_rank() * cuda_graph_batch
|
||||
|
||||
@@ -13,21 +13,20 @@ from sglang.srt.state_capturer.routed_experts import (
|
||||
extract_routed_experts_from_meta_info,
|
||||
)
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=185, suite="stage-b-test-2-gpu-large")
|
||||
register_amd_ci(
|
||||
est_time=200,
|
||||
suite="stage-b-test-2-gpu-large-amd",
|
||||
disabled="TP=2 DP=2 routed expert mismatch >15% on AMD; needs TP/DP tuning + concurrency reduction",
|
||||
)
|
||||
register_cuda_ci(est_time=400, suite="stage-c-test-4-gpu-h100")
|
||||
|
||||
# FP8 variant of Qwen3-30B-A3B: required because DeepEP normal/LL fast paths in
|
||||
# ep_moe/layer.py only run for {Fp8Config (via deep_gemm), W4AFp8Config, aiter,
|
||||
# NPU, modelopt_fp4+cutedsl}. Bf16 hits an `assert False, "deprecated"` today.
|
||||
MODEL_PATH = "Qwen/Qwen3-30B-A3B-FP8"
|
||||
|
||||
SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered"
|
||||
SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json"
|
||||
@@ -35,34 +34,43 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TestReturnRoutedExperts(CustomTestCase):
|
||||
# modified from test_hicache.py
|
||||
"""End-to-end check that --enable-return-routed-experts stays correct
|
||||
under DeepEP a2a + attn_tp_size > 1, across overlap/cuda-graph/radix
|
||||
optimisations.
|
||||
|
||||
Both servers run ``--tp 4 --dp 2 --enable-dp-attention --moe-a2a-backend
|
||||
deepep`` so attn_tp_size=2 and the all-gather hot path in
|
||||
RoutedExpertsCapturer.capture is hit on every step. Baseline disables
|
||||
overlap/cuda-graph/radix to give a deterministic ground truth; reference
|
||||
leaves them on. If the gather were skipping a rank or racing against the
|
||||
forward stream, the captured topk_ids would diverge between the two.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
|
||||
cls.baseline_args = [
|
||||
common = [
|
||||
"--enable-return-routed-experts",
|
||||
"--enable-deterministic-inference",
|
||||
"--tp",
|
||||
4,
|
||||
"--dp",
|
||||
2,
|
||||
"--enable-dp-attention",
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
# Force normal-mode dispatch: deepep auto routes decode through
|
||||
# low_latency mode whose buffer (num_max_dispatch_tokens_per_rank)
|
||||
# is undersized for cuda graph capture at default --cuda-graph-max-bs.
|
||||
"--deepep-mode",
|
||||
"normal",
|
||||
]
|
||||
cls.baseline_args = common + [
|
||||
"--disable-overlap-schedule",
|
||||
"--disable-cuda-graph",
|
||||
"--disable-radix-cache",
|
||||
"--tp",
|
||||
2,
|
||||
"--dp",
|
||||
2,
|
||||
"--enable-dp-attention",
|
||||
]
|
||||
cls.reference_args = [
|
||||
"--enable-return-routed-experts",
|
||||
"--enable-deterministic-inference",
|
||||
"--tp",
|
||||
2,
|
||||
"--dp",
|
||||
2,
|
||||
"--enable-dp-attention",
|
||||
]
|
||||
cls.sampling_args = {
|
||||
"temperature": 0,
|
||||
}
|
||||
cls.reference_args = common
|
||||
cls.sampling_args = {"temperature": 0}
|
||||
# prepare ShareGPT dataset
|
||||
dataset_path = download_and_cache_hf_file(SHAREGPT_REPO_ID, SHAREGPT_FILENAME)
|
||||
with open(dataset_path) as f:
|
||||
@@ -147,7 +155,7 @@ class TestReturnRoutedExperts(CustomTestCase):
|
||||
other_args,
|
||||
):
|
||||
process = popen_launch_server(
|
||||
DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
|
||||
MODEL_PATH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
|
||||
Reference in New Issue
Block a user