diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 0579f95b8..13677db1a 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -812,6 +812,7 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: out = runner.replay( self.bs, self.seq_lens_sum, self._replay_spec_info, self.seq_lens_cpu ) + self._publish_shared_read_done(step) raw_bs = self.raw_bs raw_num_tokens = self.raw_num_tokens num_logit_rows = raw_bs if self.prune_draft_extend_logits else raw_num_tokens @@ -829,6 +830,14 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: out.topk_index[:raw_bs], ) + def _publish_shared_read_done(self, step: int) -> None: + if step != self.speculative_num_steps - 1: + return + runner = self.runners[step] + read_done = runner.device_module.Event() + read_done.record() + runner.model_runner.shared_read_done_event = read_done + def clone_draft_probs(self) -> torch.Tensor: """Materialize the in-graph-written proposal q [raw_bs, num_steps, vocab] after replay; the clone must land before a later replay rewrites the @@ -989,4 +998,6 @@ class OneGraphMultiLayerEagleMultiStepDraftExtendCudaGraphRunner( out.topk_p[:raw_bs], out.topk_index[:raw_bs], ) - return self._cached[step] + result = self._cached[step] + self._publish_shared_read_done(step) + return result diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index cb8b82443..e5c21ceda 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -962,8 +962,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): @property def last_shared_read_runner(self): - # Multi-layer eagle has no draft forward, only draft extend. - return self._draft_worker.draft_runner + return self._draft_worker.draft_runner_list[-1] @property def spec_v2_attn_backends(self) -> tuple: diff --git a/test/registered/unit/spec/test_multi_layer_eagle_shared_read_event.py b/test/registered/unit/spec/test_multi_layer_eagle_shared_read_event.py new file mode 100644 index 000000000..18af593f8 --- /dev/null +++ b/test/registered/unit/spec/test_multi_layer_eagle_shared_read_event.py @@ -0,0 +1,172 @@ +"""Unit tests for the multi-layer EAGLE final shared-read event.""" + +import unittest +from types import SimpleNamespace + +from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import ( + MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, + OneGraphMultiLayerEagleMultiStepDraftExtendCudaGraphRunner, +) +from sglang.srt.speculative.multi_layer_eagle_worker_v2 import ( + MultiLayerEagleDraftWorker, + MultiLayerEagleWorkerV2, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +RAW_BS = 2 +RAW_NUM_TOKENS = 3 +PADDED_BS = 6 + + +class _StubEvent: + def __init__(self): + self.record_count = 0 + + def record(self): + self.record_count += 1 + + +class _StubDeviceModule: + def __init__(self): + self.events = [] + + def Event(self): + event = _StubEvent() + self.events.append(event) + return event + + +def _make_worker(num_steps): + draft_worker = MultiLayerEagleDraftWorker.__new__(MultiLayerEagleDraftWorker) + draft_worker.draft_runner_list = [ + SimpleNamespace(shared_read_done_event=None, device_timer=None) + for _ in range(num_steps) + ] + worker = MultiLayerEagleWorkerV2.__new__(MultiLayerEagleWorkerV2) + worker._draft_worker = draft_worker + return worker + + +def _make_step_runner(draft_worker, step): + out = SimpleNamespace( + next_token_logits=list(range(PADDED_BS)), + hidden_states=list(range(PADDED_BS)), + topk_p=list(range(PADDED_BS)), + topk_index=list(range(PADDED_BS)), + ) + runner = SimpleNamespace( + step=step, + raw_bs=None, + out=out, + device_module=_StubDeviceModule(), + model_runner=draft_worker.mtp_model_runner(step), + deepep_adapter=SimpleNamespace(replay=lambda: None), + ) + runner.replay = lambda bs, seq_lens_sum, spec_info, seq_lens_cpu: out + return runner + + +def _prepare_common(composite, num_steps): + worker = _make_worker(num_steps) + composite.runners = [ + _make_step_runner(worker._draft_worker, step) for step in range(num_steps) + ] + composite.speculative_num_steps = num_steps + composite.raw_bs = RAW_BS + composite.bs = RAW_BS + composite.raw_num_tokens = RAW_NUM_TOKENS + composite.seq_lens_sum = RAW_BS + composite.seq_lens_cpu = None + composite._replay_spec_info = object() + return worker + + +def _make_per_step_composite(num_steps): + composite = MultiLayerEagleMultiStepDraftExtendCudaGraphRunner.__new__( + MultiLayerEagleMultiStepDraftExtendCudaGraphRunner + ) + worker = _prepare_common(composite, num_steps) + composite.prune_draft_extend_logits = True + return composite, worker + + +def _make_one_graph_composite(num_steps): + composite = OneGraphMultiLayerEagleMultiStepDraftExtendCudaGraphRunner.__new__( + OneGraphMultiLayerEagleMultiStepDraftExtendCudaGraphRunner + ) + worker = _prepare_common(composite, num_steps) + outs = [runner.out for runner in composite.runners] + first = composite.runners[0] + first._make_graph_key = lambda bs: ("shape_key", bs) + first.backend = SimpleNamespace(replay=lambda shape_key, spec_info: outs) + return composite, worker + + +class TestLastSharedReadRunner(CustomTestCase): + def test_last_shared_read_runner_is_final_draft_runner(self): + """The barrier must read the runner the final draft-extend step writes; + pointing it at the first one leaves the event unset forever.""" + worker = _make_worker(num_steps=3) + draft_runners = worker._draft_worker.draft_runner_list + + self.assertIs(worker.last_shared_read_runner, draft_runners[-1]) + + +class TestSharedReadEventPublish(CustomTestCase): + def _assert_published(self, composite, worker, step): + runner = composite.runners[step] + events = runner.device_module.events + self.assertEqual(len(events), 1) + self.assertEqual(events[0].record_count, 1) + self.assertIs(worker.last_shared_read_runner, runner.model_runner) + self.assertIs(worker.last_shared_read_runner.shared_read_done_event, events[0]) + + def _assert_not_published(self, composite, step): + runner = composite.runners[step] + self.assertEqual(runner.device_module.events, []) + self.assertIsNone(runner.model_runner.shared_read_done_event) + + def test_per_step_replay_publishes_only_on_final_step(self): + composite, worker = _make_per_step_composite(num_steps=2) + + composite.replay(0) + self._assert_not_published(composite, 0) + self._assert_not_published(composite, 1) + + composite.replay(1) + self._assert_not_published(composite, 0) + self._assert_published(composite, worker, 1) + + def test_one_graph_replay_publishes_only_on_final_step(self): + composite, worker = _make_one_graph_composite(num_steps=2) + + composite.replay(0) + self._assert_not_published(composite, 0) + self._assert_not_published(composite, 1) + + composite.replay(1) + self._assert_not_published(composite, 0) + self._assert_published(composite, worker, 1) + + def test_single_step_replay_publishes_on_step_zero(self): + composite, worker = _make_per_step_composite(num_steps=1) + + composite.replay(0) + self._assert_published(composite, worker, 0) + + def test_replay_returns_batch_sliced_outputs(self): + composite, _ = _make_per_step_composite(num_steps=2) + + logits_output, topk_p, topk_index = composite.replay(1) + + self.assertEqual(logits_output.next_token_logits, list(range(RAW_BS))) + self.assertEqual(logits_output.hidden_states, list(range(RAW_NUM_TOKENS))) + self.assertEqual(topk_p, list(range(RAW_BS))) + self.assertEqual(topk_index, list(range(RAW_BS))) + + +if __name__ == "__main__": + unittest.main()