diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py index d07636bdf..b5aa35136 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/python/sglang/jit_kernel/dsv4/compress.py @@ -10,6 +10,7 @@ from sglang.jit_kernel.utils import ( load_jit, make_cpp_args, ) +from sglang.srt.utils import is_hip from .utils import make_name @@ -213,6 +214,20 @@ class CompressorPrefillPlan(NamedTuple): dtype=torch.uint8, pin_memory=not is_gpu_input, ) + # DP-safe empty-batch guard: a TBO ubatch (or tail batch) can have 0 + # query tokens (num_q_tokens==0) on THIS rank while other DP ranks are + # non-empty. The global TBO decision must stay uniform across ranks, so + # return an empty plan here (downstream compressor then processes 0 + # tokens = no-op) instead of skipping TBO per-rank. Avoids the + # c_plan.cuh RuntimeCheck(batch_size <= num_q_tokens) failure at B>=1. + if int(num_q_tokens) == 0 and is_hip(): + _dev = req_to_token.device + return CompressorPrefillPlan( + compress_ratio, + torch.empty((0, 16), dtype=torch.uint8, device=_dev), + torch.empty((0, 8), dtype=torch.uint8, device=_dev), + pin_buffer, + ) module = _jit_compress_plan_module() plan_c, plan_w = module.plan_prefill( req_pool_indices, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 8969a3d8f..d17daba3b 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -2148,7 +2148,9 @@ class DeepseekV4Model(nn.Module): and forward_batch.can_run_tbo and forward_batch.tbo_children is not None and forward_batch.global_forward_mode is not None - and forward_batch.global_forward_mode.is_extend() + # MTP target-verify also reports is_extend(); only real prefill + # should enter the prefill TBO strategy. + and forward_batch.global_forward_mode.is_extend_without_speculative() and not dsa_use_prefill_cp(forward_batch) and self.pp_group.world_size == 1 ) diff --git a/test/registered/amd/test_deepseek_v4_pro_fp4_tbo_mtp.py b/test/registered/amd/test_deepseek_v4_pro_fp4_tbo_mtp.py new file mode 100644 index 000000000..90713f77d --- /dev/null +++ b/test/registered/amd/test_deepseek_v4_pro_fp4_tbo_mtp.py @@ -0,0 +1,149 @@ +"""MI35x DeepSeek-V4-Pro FP4 + non-EP DP two-batch-overlap (TBO) + MTP test (8-GPU) + +End-to-end accuracy test for DeepSeek-V4-Pro (1.6T) FP4 with the non-EP DP +two-batch-overlap path on MI35x ROCm 7.2. + +Registry: nightly-amd-8-gpu-mi35x-deepseek-v4-pro-mtp suite +""" + +import os +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) + +register_amd_ci( + est_time=14400, suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro-mtp", nightly=True +) + +DEEPSEEK_V4_PRO_FP4_MODEL_PATH = os.environ.get( + "DEEPSEEK_V4_PRO_MODEL_PATH_FP4", "deepseek-ai/DeepSeek-V4-Pro" +) +# Pro is 1.6T; weight load + warmup is much longer than Flash 285B. +SERVER_LAUNCH_TIMEOUT = 5400 +FLASHMLA_BACKEND = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "unified_kv_triton") + + +COMMON_ENV_VARS = { + "SGLANG_DEFAULT_THINKING": "1", + "SGLANG_DSV4_REASONING_EFFORT": "max", + "SGLANG_USE_ROCM700A": "0", + "SGLANG_HACK_FLASHMLA_BACKEND": FLASHMLA_BACKEND, + "AITER_BF16_FP8_MOE_BOUND": "0", + # DP TP-MoE collective path that non-EP DP TBO overlaps. + "SGLANG_DP_USE_GATHERV": "1", + "SGLANG_DP_USE_REDUCE_SCATTER": "1", + "SGLANG_SHARED_EXPERT_TP1": "1", + "SGLANG_DP_SHARED_EXPERT_LOCAL": "1", + # ROCm HSA-resource stability for TBO at high concurrency. + "GPU_MAX_HW_QUEUES": "5", +} + +# FP4 variant +FP4_ENV_VARS = { + "SGLANG_DSV4_FP4_EXPERTS": "true", +} + + +class TestDeepseekV4ProFp4TboMTP(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = DEEPSEEK_V4_PRO_FP4_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + + env = os.environ.copy() + env.update(COMMON_ENV_VARS) + env.update(FP4_ENV_VARS) + + other_args = [ + "--trust-remote-code", + "--tp", + "8", + # DP attention + TBO: non-EP DP TP-MoE two-batch-overlap. DP TBO is + # selected because moe_a2a_backend stays 'none'; no opt-in env needed. + "--dp", + "8", + "--enable-dp-attention", + "--enable-prefill-delayer", + "--enable-two-batch-overlap", + "--disable-radix-cache", + "--attention-backend", + "dsv4", + "--kv-cache-dtype", + "fp8_e4m3", + # MTP / EAGLE speculative decoding (NextN head from the base model). + # With spec on, target-verify batches also report is_extend(); the + # prefill TBO gate is is_extend_without_speculative() (PR #30238). + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + "--max-running-requests", + "512", + "--cuda-graph-max-bs", + "512", + "--page-size", + "256", + "--mem-fraction-static", + "0.90", + "--swa-full-tokens-ratio", + "0.15", + # global chunk; DP-attention divides by dp_size=8 -> 8192/rank. + "--chunked-prefill-size", + "65536", + "--disable-shared-experts-fusion", + "--tool-call-parser", + "deepseekv4", + "--reasoning-parser", + "deepseek-v4", + ] + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + env=env, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k_tbo_mtp(self): + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=1319, + parallel=512, + 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=}") + + if is_in_ci(): + write_github_step_summary( + f"### test_gsm8k (deepseek-v4-pro-fp4 DP+TBO+MTP, {FLASHMLA_BACKEND})\n" + f'{metrics["accuracy"]=:.3f}\n' + ) + # TBO+MTP must not regress accuracy vs the non-TBO MTP baseline. + self.assertGreater(metrics["accuracy"], 0.91) + + +if __name__ == "__main__": + unittest.main()