From fd6f96bf96f1402597810a24c78ff32071dcd1c5 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Thu, 17 Sep 2026 14:33:25 +0800 Subject: [PATCH] Re-land dp-attention local control broadcast test (#39835) Co-authored-by: Sam Shleifer --- .../sglang/test/kits/pause_generation_kit.py | 2 + ...st_dp_attention_local_control_broadcast.py | 128 ++++++++++++++++++ 2 files changed, 130 insertions(+) create mode 100644 test/registered/e2e/dp_attn/test_dp_attention_local_control_broadcast.py diff --git a/python/sglang/test/kits/pause_generation_kit.py b/python/sglang/test/kits/pause_generation_kit.py index 00c1ef4cf..dac85a73e 100644 --- a/python/sglang/test/kits/pause_generation_kit.py +++ b/python/sglang/test/kits/pause_generation_kit.py @@ -20,6 +20,7 @@ class PauseResumeInPlaceMixin: pause_num_requests: int = 32 pause_max_new_tokens: int = 512 pause_duration: float = 5 + pause_ignore_eos: bool = False pause_generate_url: str = "" pause_target_urls: list = [] @@ -36,6 +37,7 @@ class PauseResumeInPlaceMixin: "sampling_params": { "temperature": 0.8, "max_new_tokens": self.pause_max_new_tokens, + "ignore_eos": self.pause_ignore_eos, }, }, timeout=_REQUEST_TIMEOUT, diff --git a/test/registered/e2e/dp_attn/test_dp_attention_local_control_broadcast.py b/test/registered/e2e/dp_attn/test_dp_attention_local_control_broadcast.py new file mode 100644 index 000000000..556d986d3 --- /dev/null +++ b/test/registered/e2e/dp_attn/test_dp_attention_local_control_broadcast.py @@ -0,0 +1,128 @@ +"""Control-message routing under DP attention with local control broadcast. + +With --enable-dp-attention-local-control-broadcast the tokenizer's control +messages (flush_cache, pause_generation, continue_generation) go to every DP +group leader over ZMQ and are broadcast within attn_tp_group instead of the +full tp_group. A rank that misses one of them either keeps generating while +the others are paused or hangs the next gloo broadcast, so this test drives +the whole tokenizer -> DP controller -> scheduler path on a tp4/dp2 server +(attn_tp_size=2, two DP groups of two ranks each) and checks that greedy +outputs survive flush / pause / continue round-trips. +""" + +import time +import unittest +from concurrent.futures import ThreadPoolExecutor + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.pause_generation_kit import PauseResumeInPlaceMixin +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST_MLA, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=400, stage="base-c", runner_config="4-gpu-h100") + +_INPUT_IDS = list(range(10, 30)) +_REQUEST_TIMEOUT = 120 + + +class TestDPAttentionLocalControlBroadcast(PauseResumeInPlaceMixin, CustomTestCase): + # This checkpoint can emit EOS on the first token. + pause_ignore_eos = True + + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp", + "4", + "--dp", + "2", + "--enable-dp-attention", + "--enable-dp-attention-local-control-broadcast", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def _generate_greedy(self, n: int, max_new_tokens: int = 16) -> list[list[int]]: + resp = requests.post( + f"{self.base_url}/generate", + json={ + "input_ids": [_INPUT_IDS] * n, + "sampling_params": { + "temperature": 0, + "max_new_tokens": max_new_tokens, + }, + }, + timeout=_REQUEST_TIMEOUT, + ) + resp.raise_for_status() + outputs = resp.json() + if isinstance(outputs, dict): + outputs = [outputs] + return [o["output_ids"] for o in outputs] + + def _post_control(self, route: str, payload: dict): + requests.post( + f"{self.base_url}/{route}", json=payload, timeout=60 + ).raise_for_status() + + def test_flush_and_pause_round_trips_keep_greedy_outputs(self): + baseline = self._generate_greedy(n=4) + self.assertEqual(len(baseline), 4) + for ids in baseline: + self.assertGreater(len(ids), 0) + + self._post_control("flush_cache", {}) + self.assertEqual( + self._generate_greedy(n=4), + baseline, + "greedy outputs changed after flush_cache", + ) + + self._post_control("pause_generation", {"mode": "in_place"}) + with ThreadPoolExecutor(max_workers=1) as pool: + blocked = pool.submit(self._generate_greedy, 1, 5) + time.sleep(2.0) + self.assertFalse( + blocked.done(), + "generation completed while paused: at least one rank did not " + "receive the pause control message", + ) + self._post_control("continue_generation", {}) + resumed = blocked.result(timeout=60) + self.assertEqual(len(resumed), 1) + self.assertGreater(len(resumed[0]), 0) + + self.assertEqual( + self._generate_greedy(n=4), + baseline, + "greedy outputs changed after pause/continue", + ) + + self._post_control("flush_cache", {}) + self.assertEqual( + self._generate_greedy(n=2), + baseline[:2], + "greedy outputs changed after the second flush_cache", + ) + + +if __name__ == "__main__": + unittest.main()