fix(spec_v2): support EAGLE with spec_v2

1. add pass for spec_v2 in base_attn

2. fix: EAGLE with spec_v2 overlap Grammar accept_token failed

Signed-off-by: wxiwnd <wxiwnd@outlook.com>

python/sglang/srt/disaggregation/decode.py

@@ -137,9 +137,9 @@ class DecodeReqToTokenPool:
         # Indices of reqs that already have a req_pool_idx and will reuse
         # their existing slot (e.g. chunked prefill continuing across chunks).
         reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]
-        assert (
-            len(reusing) <= 1
-        ), "only one chunked request may reuse req_pool_idx in a batch"
+        assert len(reusing) <= 1, (
+            "only one chunked request may reuse req_pool_idx in a batch"
+        )
         assert all(
             reqs[i].is_chunked > 0 or reqs[i].kv_committed_len > 0 for i in reusing
         ), "reusing request must be chunked or have committed KV"
@@ -166,7 +166,6 @@ class DecodeReqToTokenPool:
 
 
 class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
-
     def __init__(
         self,
         size: int,
@@ -793,9 +792,9 @@ class DecodePreallocQueue:
         """Pre-allocate the memory for req_to_token and token_kv_pool"""
         req_pool_indices = self.req_to_token_pool.alloc([req])
 
-        assert (
-            req_pool_indices is not None
-        ), "req_pool_indices is full! There is a bug in memory estimation."
+        assert req_pool_indices is not None, (
+            "req_pool_indices is full! There is a bug in memory estimation."
+        )
 
         # Alloc all tokens for the prebuilt req (except for the reserved input token for decoding)
         fill_len = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
@@ -814,9 +813,9 @@ class DecodePreallocQueue:
                 extend_num_tokens=fill_len,
             )
 
-        assert (
-            kv_loc is not None
-        ), "KV cache is full! There is a bug in memory estimation."
+        assert kv_loc is not None, (
+            "KV cache is full! There is a bug in memory estimation."
+        )
 
         self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
 
@@ -1008,7 +1007,6 @@ class DecodeTransferQueue:
 
 
 class SchedulerDisaggregationDecodeMixin:
-
     @torch.no_grad()
     def event_loop_normal_disagg_decode(self: Scheduler):
         """A normal scheduler loop for decode worker in disaggregation mode."""
@@ -1023,6 +1021,18 @@ class SchedulerDisaggregationDecodeMixin:
             # Get the next batch to run
             batch = self.get_next_disagg_decode_batch_to_run()
             self.cur_batch = batch
+            disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch)
+
+            def pop_and_process():
+                # Process the results of the last batch
+                tmp_batch, tmp_result = self.result_queue.popleft()
+                self.process_batch_result(tmp_batch, tmp_result)
+
+            # If we need grammar sync (spec + grammar), process the last batch
+            # results first so that the grammar state is up-to-date before
+            # generating the bitmask for the current batch.
+            if disable_overlap_for_batch:
+                pop_and_process()
 
             # Launch the current batch
             if batch:
@@ -1040,6 +1050,11 @@ class SchedulerDisaggregationDecodeMixin:
         self.result_queue = deque()
         self.last_batch: Optional[ScheduleBatch] = None
 
+        def pop_and_process():
+            # Process the results of the last batch
+            tmp_batch, tmp_result = self.result_queue.popleft()
+            self.process_batch_result(tmp_batch, tmp_result)
+
         while True:
             # Receive requests
             recv_reqs = self.recv_requests()
@@ -1050,6 +1065,13 @@ class SchedulerDisaggregationDecodeMixin:
             # Get the next batch to run
             batch = self.get_next_disagg_decode_batch_to_run()
             self.cur_batch = batch
+            disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch)
+
+            # If we need grammar sync (spec + grammar), process the last batch
+            # results first so that the grammar state is up-to-date before
+            # generating the bitmask for the current batch.
+            if self.last_batch and disable_overlap_for_batch:
+                pop_and_process()
 
             # Launch the current batch
             if batch:
@@ -1060,8 +1082,8 @@ class SchedulerDisaggregationDecodeMixin:
 
             # Process the last batch
             if self.last_batch:
-                tmp_batch, tmp_result = self.result_queue.popleft()
-                self.process_batch_result(tmp_batch, tmp_result)
+                if not disable_overlap_for_batch:
+                    pop_and_process()
             elif batch is None:
                 self.self_check_during_idle()
 

python/sglang/srt/layers/attention/base_attn_backend.py

@@ -75,7 +75,10 @@ class AttentionBackend(ABC):
         Here, we need to redo the computation of all metadata of the attention backend
         that depends on tree mask and position buffers.
         """
-        raise NotImplementedError()
+        # if self.topk <= 1:
+        #     return
+        # raise NotImplementedError()
+        pass
 
     @debug_kernel_api
     def forward(

test/registered/unit/disaggregation/test_overlap_disagg_decode_event_loop.py

@@ -0,0 +1,85 @@
+"""Unit tests for overlap disaggregation decode event loop ordering."""
+
+import unittest
+from types import SimpleNamespace
+from typing import Any, cast
+
+from sglang.srt.disaggregation.decode import SchedulerDisaggregationDecodeMixin
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=6, suite="stage-a-test-cpu")
+
+
+class FakeBatch:
+    def __init__(self, label: str):
+        self.label = label
+
+    def copy(self):
+        return self
+
+
+class TestOverlapDisaggDecodeEventLoop(CustomTestCase):
+    def test_processes_previous_result_before_disabled_overlap_launch(self):
+        class StopLoop(Exception):
+            pass
+
+        batch1 = FakeBatch("batch-1")
+        batch2 = FakeBatch("batch-2")
+        batches = iter([batch1, batch2])
+        events = []
+
+        def get_next_batch():
+            return next(batches)
+
+        def run_batch(batch):
+            result = f"result-{batch.label}"
+            events.append(f"run:{batch.label}")
+            return result
+
+        def process_batch_result(batch, result):
+            events.append(f"process:{batch.label}")
+
+        def launch_batch_sample_if_needed(result):
+            events.append(f"sample:{result}")
+            if result == "result-batch-2":
+                raise StopLoop()
+
+        scheduler = cast(
+            Any,
+            SimpleNamespace(
+                recv_requests=lambda: [],
+                process_input_requests=lambda reqs: None,
+                process_decode_queue=lambda: None,
+                get_next_disagg_decode_batch_to_run=get_next_batch,
+                run_batch=run_batch,
+                process_batch_result=process_batch_result,
+                is_disable_overlap_for_batch=lambda batch: batch is batch2,
+                launch_batch_sample_if_needed=launch_batch_sample_if_needed,
+                self_check_during_idle=lambda: None,
+            ),
+        )
+
+        with self.assertRaises(StopLoop):
+            SchedulerDisaggregationDecodeMixin.event_loop_overlap_disagg_decode(
+                scheduler
+            )
+
+        self.assertEqual(
+            events,
+            [
+                "run:batch-1",
+                "sample:result-batch-1",
+                "process:batch-1",
+                "run:batch-2",
+                "sample:result-batch-2",
+            ],
+        )
+        self.assertEqual(len(scheduler.result_queue), 1)
+        queued_batch, queued_result = scheduler.result_queue[0]
+        self.assertIs(queued_batch, batch2)
+        self.assertEqual(queued_result, "result-batch-2")
+
+
+if __name__ == "__main__":
+    unittest.main()