From 350fbbf4dc1eb9c1b6fd0baadd1caf24ed935647 Mon Sep 17 00:00:00 2001 From: Jincong Chen Date: Mon, 22 Dec 2025 05:16:46 +0800 Subject: [PATCH] fix ds3.2 nsa backend prefill TBO (#14901) --- .../srt/layers/attention/tbo_backend.py | 3 + python/sglang/srt/models/deepseek_v2.py | 9 ++- python/sglang/srt/server_args.py | 9 +++ test/srt/ep/test_deepep_large.py | 55 +++++++++++++++++++ test/srt/run_suite.py | 2 +- 5 files changed, 76 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index 04e7b07c7..494d82d80 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -185,6 +185,9 @@ class TboAttnBackend(AttentionBackend): def forward_decode(self, *args, **kwargs): return self.primary.forward_decode(*args, **kwargs) + def get_indexer_metadata(self, layer_id: int, forward_batch: "ForwardBatch"): + return self.primary.get_indexer_metadata(layer_id, forward_batch) + def _init_forward_metadata_cuda_graph_split( fn_name: str, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 7e5153f6e..4fd8376e4 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -67,6 +67,7 @@ from sglang.srt.layers.attention.nsa.utils import ( is_nsa_enable_prefill_cp, prepare_input_dp_with_cp_dsa, ) +from sglang.srt.layers.attention.tbo_backend import TboAttnBackend from sglang.srt.layers.attention.utils import concat_and_cast_mha_k_triton from sglang.srt.layers.communicator import ( LayerCommunicator, @@ -425,7 +426,10 @@ def handle_attention_nsa(attn, forward_batch): Dispatch logic is centralized in NativeSparseAttnBackend.set_nsa_prefill_impl and executed in init_forward_metadata. Read the decision from backend.use_mha. """ + backend = forward_batch.attn_backend + if isinstance(backend, TboAttnBackend): # if enable tbo, get primary backend + backend = backend.primary if hasattr(backend, "use_mha") and backend.use_mha: return AttnForwardMethod.MHA_ONE_SHOT return AttnForwardMethod.MLA @@ -2670,7 +2674,10 @@ class DeepseekV2AttentionMLA(nn.Module): Returns: (kv_a, k_pe) both in BF16 """ - kv_indices = forward_batch.attn_backend.forward_metadata.page_table_1_flattened + backend = forward_batch.attn_backend + if isinstance(backend, TboAttnBackend): # if enable tbo, get primary backend + backend = backend.primary + kv_indices = backend.forward_metadata.page_table_1_flattened assert ( kv_indices is not None ), "page_table_1_flattened should have been generated for FP8 MHA path" diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 15bde35eb..3eeabca04 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -727,6 +727,9 @@ class ServerArgs: # Handle any other necessary validations. self._handle_other_validations() + # Handle two-batch overlap settings. + self._handle_two_batch_overlap() + def _handle_deprecated_args(self): # Handle deprecated tool call parsers deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"} @@ -2392,6 +2395,12 @@ class ServerArgs: self.preferred_sampling_params ) + def _handle_two_batch_overlap(self): + if self.enable_two_batch_overlap and self.moe_a2a_backend == "none": + raise ValueError( + "When enabling two batch overlap, moe_a2a_backend cannot be 'none'." + ) + @staticmethod def add_cli_args(parser: argparse.ArgumentParser): diff --git a/test/srt/ep/test_deepep_large.py b/test/srt/ep/test_deepep_large.py index 1c4a91e6a..a98999376 100644 --- a/test/srt/ep/test_deepep_large.py +++ b/test/srt/ep/test_deepep_large.py @@ -5,6 +5,7 @@ import requests from sglang.srt.utils import kill_process_tree from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.send_one import BenchArgs, send_one_prompt from sglang.test.test_utils import ( DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -13,6 +14,8 @@ from sglang.test.test_utils import ( popen_launch_server, ) +DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp" + class TestDeepseek(CustomTestCase): @classmethod @@ -153,5 +156,57 @@ class TestDeepseekMTP(CustomTestCase): self.assertGreater(avg_spec_accept_length, 1.85) +class TestDeepseekV32TBO(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEEPSEEK_V32_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + other_args = [ + "--trust-remote-code", + "--tp", + "8", + "--dp", + "8", + "--enable-dp-attention", + "--enable-two-batch-overlap", + "--moe-a2a-backend", + "deepep", + "--cuda-graph-max-bs", + "32", + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_a_gsm8k( + self, + ): # Append an "a" to make this test run first (alphabetically) to warm up the server + args = SimpleNamespace( + num_shots=20, + data_path=None, + num_questions=1400, + parallel=1400, + max_new_tokens=512, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"{metrics=}") + self.assertGreater(metrics["accuracy"], 0.92) + + def test_bs_1_speed(self): + args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048) + acc_length, speed = send_one_prompt(args) + + print(f"{speed=:.2f}") + + if __name__ == "__main__": unittest.main() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 8ff79a123..c27e890a5 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -186,7 +186,7 @@ suites = { # TestFile("ep/test_mooncake_ep_small.py", 450), ], "per-commit-8-gpu-h200-deepep": [ - TestFile("ep/test_deepep_large.py", 338), + TestFile("ep/test_deepep_large.py", 563), ], "quantization_test": [ TestFile("quant/test_awq.py", 163),