Prevent spec-v2 decode warmup races

Port the fix-decode spec-v2 ownership and plan-stream fixes onto the current branch. Draft extend now keeps the scheduler-owned batch lengths committed until acceptance, binds the draft runner to the draft-extend attention backend, keeps speculative KV allocation monotonic, and lets non-graph target verify initialize metadata after DP padding. The worker also records rebound tensors on the forward stream and orders plan-stream metadata work after current-stream inputs are available.

Constraint: fix-decode commits cd8e47ed9c and 60e3956d9c address CUDA illegal-address failures in spec-v2 decode warmup paths.

Rejected: Cherry-pick the commits blindly | the current branch has intervening decode changes, so a minimal manual port kept the patch surface to the affected files.

Confidence: medium

Scope-risk: moderate

Directive: Do not move target-verify non-graph metadata initialization back into prepare_for_v2_verify without validating DP padding and NSA metadata ordering.

Tested: RED local pytest test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py -q failed 8/8 before production changes.

Tested: PYTHONPATH=python python3 -m pytest test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py -q passed 8/8 locally.

Tested: PYTHONPATH=python python3 -m py_compile python/sglang/srt/speculative/eagle_info_v2.py python/sglang/srt/speculative/eagle_worker_v2.py python/sglang/srt/speculative/spec_utils.py passed locally.

Tested: Remote g0034:cjy-glm5-new pytest for test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py passed 8/8, plus remote py_compile for the three speculative modules.

Tested: Local and remote sha256 sums matched for all four synced files.

Not-tested: Full spec-v2 decode server restart/warmup under production traffic.

python/sglang/srt/speculative/eagle_info_v2.py

@@ -93,14 +93,17 @@ class EagleDraftInputV2Mixin:
         cur_kv_lens_cpu = []
         nxt_kv_lens_cpu = []
         num_needed_tokens = 0
-        alloc_len_per_decode = get_alloc_len_per_decode()
+        reserve_len_per_decode = 2 * get_alloc_len_per_decode()
         for r in batch.reqs:
-            # Over-allocation happens here
-            x = r.kv_committed_len + 2 * alloc_len_per_decode - r.kv_allocated_len
-            cur_kv_lens_cpu.append(r.kv_allocated_len)
-            nxt_kv_lens_cpu.append(r.kv_allocated_len + x)
-            num_needed_tokens += x
-            r.kv_allocated_len += x
+            # Over-allocation happens here.  In overlap mode kv_committed_len can
+            # lag behind kv_allocated_len by a previous speculative reserve; the
+            # allocation watermark is monotonic and must never be shrunk here.
+            cur = r.kv_allocated_len
+            nxt = max(cur, r.kv_committed_len + reserve_len_per_decode)
+            cur_kv_lens_cpu.append(cur)
+            nxt_kv_lens_cpu.append(nxt)
+            num_needed_tokens += nxt - cur
+            r.kv_allocated_len = nxt
             r.decode_batch_idx += 1
 
         cur_kv_lens_cpu = torch.tensor(cur_kv_lens_cpu, dtype=torch.int32, device="cpu")
@@ -190,9 +193,6 @@ class EagleDraftInputV2Mixin:
 
         batch.spec_info = self
         batch.input_ids = predict
-        batch.seq_lens = batch.seq_lens + num_draft_tokens
-        batch.seq_lens_cpu = batch.seq_lens_cpu + num_draft_tokens
-        batch.seq_lens_sum += extend_num_tokens
         batch.extend_seq_lens = [num_draft_tokens for _ in range(len(batch.seq_lens))]
         batch.extend_prefix_lens = seq_lens_cpu_.tolist()
         batch.extend_num_tokens = extend_num_tokens
@@ -203,6 +203,13 @@ class EagleDraftInputV2Mixin:
             else ForwardMode.DRAFT_EXTEND_V2
         )
         forward_batch = ForwardBatch.init_new(batch, draft_model_runner)
+        # Draft extend writes num_draft_tokens future slots.  The attention
+        # metadata for this forward sees post-write lengths, but the shared
+        # ModelWorkerBatch must remain at the pre-draft committed lengths.  The
+        # scheduler/verify path advances the source batch only after acceptance.
+        forward_batch.seq_lens = forward_batch.seq_lens + num_draft_tokens
+        forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu + num_draft_tokens
+        forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum().item())
         can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
         if not batch.forward_mode.is_idle() and not can_cuda_graph:
             draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
@@ -261,11 +268,10 @@ class EagleVerifyInputV2Mixin:
         )
         if can_run_cuda_graph:
             target_worker.model_runner.graph_runner.replay_prepare(verify_forward_batch)
-        else:
-            if not batch.forward_mode.is_idle():
-                target_worker.model_runner.attn_backend.init_forward_metadata(
-                    verify_forward_batch
-                )
+        # Non-cuda-graph target verify must initialize attention metadata inside
+        # ModelRunner.forward_extend, after prepare_mlp_sync_batch has applied DP
+        # padding.  Planning here uses pre-pad shapes and can corrupt NSA/DSA
+        # indexer metadata.
 
         return verify_forward_batch, can_run_cuda_graph
 

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -51,6 +51,8 @@ from sglang.srt.speculative.spec_utils import (
     load_token_map,
     maybe_detect_nan,
     maybe_detect_oob,
+    record_stream_each,
+    record_stream_for_v2_verify,
     select_top_k_tokens,
 )
 from sglang.srt.utils.common import (
@@ -246,6 +248,8 @@ class EagleDraftWorker(BaseDraftWorker):
         )
 
         self.draft_runner.draft_attn_backend = self.draft_attn_backend
+        if self.draft_extend_attn_backend is not None:
+            self.draft_runner.attn_backend = self.draft_extend_attn_backend
         self.tree_mask_mode = TreeMaskMode.FULL_MASK
 
     def init_cuda_graphs(self):
@@ -552,6 +556,7 @@ class EagleDraftWorker(BaseDraftWorker):
     def _draft_extend_for_decode(
         self, batch: ModelWorkerBatch, batch_result: GenerationBatchResult
     ):
+        fwd_stream = torch.get_device_module(self.device).current_stream()
         # Batch 2: Draft extend
         draft_input = EagleDraftInput(
             hidden_states=batch_result.logits_output.hidden_states,
@@ -565,11 +570,18 @@ class EagleDraftWorker(BaseDraftWorker):
             - 1
         )
 
+        # Cast before entering the plan stream.  Doing dtype conversion inside
+        # the planning stream creates an implicit cross-stream dependency for
+        # the draft-extend metadata path.
+        next_token_ids = batch_result.next_token_ids.to(torch.int64)
+
         # Prepare for draft extend in a separate stream
+        if self.plan_stream:
+            self.plan_stream.wait_stream(fwd_stream)
         with self.plan_stream_ctx:
             forward_batch = draft_input.prepare_for_extend_to_fill_draft_kvcache(
                 batch,
-                batch_result.next_token_ids,
+                next_token_ids,
                 self.speculative_num_draft_tokens,
                 self.draft_runner,
                 self.cuda_graph_runner_for_draft_extend,
@@ -822,20 +834,18 @@ class EAGLEWorkerV2(BaseSpecWorker):
             return batch_output
 
     def verify(self, batch: ModelWorkerBatch):
-        # Since batch.seq_lens is allocated in another stream, we need
-        # record_stream() to prevent pytorch gc and reuse the gpu memory
-        # while forward_stream is still running.
-        batch.seq_lens.record_stream(
-            torch.get_device_module(self.device).current_stream()
-        )
+        fwd_stream = torch.get_device_module(self.device).current_stream()
 
         # Parse args
         verify_input: EagleVerifyInput = batch.spec_info
+        record_stream_for_v2_verify(batch, verify_input, fwd_stream)
         verify_input.num_tokens_per_req = self.speculative_num_steps + 1
         bs = len(batch.seq_lens)
 
         # Batch 1: Target verify
         # Prepare for target verify in a separate stream
+        if self.plan_stream:
+            self.plan_stream.wait_stream(fwd_stream)
         with self.plan_stream_ctx:
             verify_forward_batch, can_run_cuda_graph = (
                 verify_input.prepare_for_v2_verify(
@@ -845,6 +855,10 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 )
             )
 
+        # Cover post-prepare rebinds: draft_token/input_ids and plan-stream
+        # allocated out_cache_loc are read by kernels on the forward stream.
+        record_stream_each((batch.input_ids, batch.out_cache_loc), fwd_stream)
+
         # Correct some buffers due to the overlap plan
         if self.plan_stream:
             torch.get_device_module(self.device).current_stream().wait_stream(
@@ -876,7 +890,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
             model_worker_batch=None,
             forward_batch=verify_forward_batch,
             is_verify=True,
-            skip_attn_backend_init=True,
+            skip_attn_backend_init=can_run_cuda_graph,
         )
         logits_output = forward_batch_output.logits_output
 

python/sglang/srt/speculative/spec_utils.py

@@ -41,6 +41,50 @@ else:
 logger = logging.getLogger(__name__)
 
 
+def record_stream_each(tensors, stream):
+    """Record CUDA tensor storage use on an additional stream.
+
+    Spec v2 rebinding can drop the only Python reference to tensors while
+    already-enqueued kernels still read them.  record_stream prevents the
+    CUDA caching allocator from recycling those blocks before ``stream`` has
+    completed.
+    """
+    for tensor in tensors:
+        if isinstance(tensor, torch.Tensor) and tensor.is_cuda:
+            tensor.record_stream(stream)
+
+
+def record_stream_for_v2_verify(batch, verify_input, fwd_stream):
+    """Protect spec-v2 verify tensors that may be rebound during planning.
+
+    ``prepare_for_v2_verify`` rebinds ``batch.input_ids`` and
+    ``batch.out_cache_loc``.  The old tensors can still be consumed by kernels
+    queued on the forward stream, so mark all pre-prepare candidates as used by
+    that stream before the rebind happens.  Callers must separately record the
+    post-prepare rebinds.
+    """
+    candidates = [
+        getattr(batch, "seq_lens", None),
+        getattr(batch, "req_pool_indices", None),
+        getattr(batch, "input_ids", None),
+        getattr(batch, "out_cache_loc", None),
+    ]
+    if verify_input is not None:
+        for attr in (
+            "draft_token",
+            "custom_mask",
+            "positions",
+            "retrieve_index",
+            "retrieve_next_token",
+            "retrieve_next_sibling",
+            "retrive_index",
+            "retrive_next_token",
+            "retrive_next_sibling",
+        ):
+            candidates.append(getattr(verify_input, attr, None))
+    record_stream_each(candidates, fwd_stream)
+
+
 # Simulate acceptance length for benchmarking purposes
 SIMULATE_ACC_LEN = envs.SGLANG_SIMULATE_ACC_LEN.get()  # turn off if < 0
 SIMULATE_ACC_METHOD = envs.SGLANG_SIMULATE_ACC_METHOD.get()

test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py

@@ -0,0 +1,232 @@
+import ast
+from pathlib import Path
+
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+
+
+def _parse_module(relative_path: str) -> ast.Module:
+    return ast.parse((REPO_ROOT / relative_path).read_text())
+
+
+def _find_function(tree: ast.AST, name: str) -> ast.FunctionDef:
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == name:
+            return node
+    raise AssertionError(f"function {name!r} not found")
+
+
+def _assigned_attrs(func: ast.FunctionDef) -> set[tuple[str, str]]:
+    assigned: set[tuple[str, str]] = set()
+    for node in ast.walk(func):
+        if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
+            targets = []
+            if isinstance(node, ast.Assign):
+                targets.extend(node.targets)
+            else:
+                targets.append(node.target)
+            for target in targets:
+                if (
+                    isinstance(target, ast.Attribute)
+                    and isinstance(target.value, ast.Name)
+                ):
+                    assigned.add((target.value.id, target.attr))
+    return assigned
+
+
+def test_eagle_v2_draft_extend_prepare_does_not_advance_source_batch_lengths():
+    """Draft-extend metadata may use post-write lengths, but the scheduler batch
+    must keep pre-draft lengths.
+
+    Mutating ModelWorkerBatch.seq_lens in the prepare phase makes the NSA page
+    table contract depend on future draft slots before the forward batch owns its
+    metadata, which can surface as async CUDA illegal-address failures.
+    """
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_info_v2.py")
+    func = _find_function(tree, "prepare_for_extend_to_fill_draft_kvcache")
+
+    assigned = _assigned_attrs(func)
+    assert ("batch", "seq_lens") not in assigned
+    assert ("batch", "seq_lens_cpu") not in assigned
+    assert ("batch", "seq_lens_sum") not in assigned
+
+    assert ("forward_batch", "seq_lens") in assigned
+    assert ("forward_batch", "seq_lens_cpu") in assigned
+    assert ("forward_batch", "seq_lens_sum") in assigned
+
+
+def test_eagle_v2_binds_draft_runner_to_draft_extend_attention_backend():
+    tree = _parse_module("python/sglang/srt/speculative/eagle_worker_v2.py")
+    func = _find_function(tree, "init_attention_backend")
+
+    for node in ast.walk(func):
+        if not isinstance(node, ast.Assign):
+            continue
+        for target in node.targets:
+            if (
+                isinstance(target, ast.Attribute)
+                and target.attr == "attn_backend"
+                and isinstance(target.value, ast.Attribute)
+                and target.value.attr == "draft_runner"
+                and isinstance(target.value.value, ast.Name)
+                and target.value.value.id == "self"
+                and isinstance(node.value, ast.Attribute)
+                and node.value.attr == "draft_extend_attn_backend"
+                and isinstance(node.value.value, ast.Name)
+                and node.value.value.id == "self"
+            ):
+                return
+
+    raise AssertionError("draft_runner.attn_backend is not bound to draft_extend backend")
+
+
+
+def _function_calls(func: ast.FunctionDef, name: str) -> int:
+    count = 0
+    for node in ast.walk(func):
+        if isinstance(node, ast.Call):
+            callee = node.func
+            if isinstance(callee, ast.Name) and callee.id == name:
+                count += 1
+            elif isinstance(callee, ast.Attribute) and callee.attr == name:
+                count += 1
+    return count
+
+
+def test_eagle_v2_verify_records_rebound_tensors_across_streams():
+    """Spec v2 verify rebinds tensors while forward kernels still use them.
+
+    The worker must record both pre-prepare tensors and post-prepare rebinds on
+    the forward stream; otherwise the CUDA caching allocator may recycle storage
+    while target verify or draft-extend metadata kernels are still reading it.
+    """
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_worker_v2.py")
+    func = _find_function(tree, "verify")
+
+    assert _function_calls(func, "record_stream_for_v2_verify") == 1
+    assert _function_calls(func, "record_stream_each") >= 1
+
+
+def test_eagle_v2_prepare_for_decode_never_shrinks_overallocated_kv():
+    """Spec-v2 decode over-allocation is monotonic per request.
+
+    In overlap scheduling, kv_committed_len can lag behind kv_allocated_len by a
+    previous speculative reserve.  A later prepare step must not compute a
+    negative delta and shrink kv_allocated_len; doing so leaves req_to_token and
+    allocator ownership out of sync and can surface as an async illegal-address
+    failure at a later CUDA sync point.
+    """
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_info_v2.py")
+    func = _find_function(tree, "prepare_for_decode")
+
+    has_monotonic_clamp = False
+    shrinks_in_place = False
+    for node in ast.walk(func):
+        if (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id == "max"
+        ):
+            has_monotonic_clamp = True
+        if (
+            isinstance(node, ast.AugAssign)
+            and isinstance(node.target, ast.Attribute)
+            and node.target.attr == "kv_allocated_len"
+        ):
+            shrinks_in_place = True
+
+    assert has_monotonic_clamp, "prepare_for_decode must clamp next KV len with max(cur, target)"
+    assert not shrinks_in_place, "prepare_for_decode must assign the clamped len, not += a possibly negative delta"
+
+
+def test_eagle_v2_target_verify_non_graph_metadata_is_initialized_post_padding():
+    """Target verify metadata must be planned after DP padding on non-graph path.
+
+    prepare_for_v2_verify runs on the plan stream before ModelRunner applies DP
+    padding.  Pre-planning NSA/DSA metadata there and then forcing
+    skip_attn_backend_init=True leaves the forward path with pre-pad metadata,
+    which can corrupt indexer/page-table kernels and surface later as an async
+    illegal memory access.
+    """
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_info_v2.py")
+    func = _find_function(tree, "prepare_for_v2_verify")
+
+    assert _function_calls(func, "init_forward_metadata") == 0
+
+
+def test_eagle_v2_target_verify_skips_forward_metadata_only_for_cuda_graph():
+    """Only cuda-graph verify has preplanned metadata.
+
+    The non-graph path must allow TpModelWorker/ModelRunner to initialize
+    metadata after prepare_mlp_sync_batch padding.  A hard-coded
+    skip_attn_backend_init=True is unsafe for DP attention.
+    """
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_worker_v2.py")
+    func = _find_function(tree, "verify")
+
+    matching_keywords = []
+    for node in ast.walk(func):
+        if not isinstance(node, ast.Call):
+            continue
+        callee = node.func
+        if not (
+            isinstance(callee, ast.Attribute)
+            and callee.attr == "forward_batch_generation"
+        ):
+            continue
+        for kw in node.keywords:
+            if kw.arg == "skip_attn_backend_init":
+                matching_keywords.append(kw.value)
+
+    assert len(matching_keywords) == 1
+    value = matching_keywords[0]
+    assert isinstance(value, ast.Name) and value.id == "can_run_cuda_graph"
+
+
+def _has_plan_stream_wait_before_context(func: ast.FunctionDef) -> bool:
+    body = list(func.body)
+    for idx, node in enumerate(body):
+        if not isinstance(node, ast.With):
+            continue
+        if not any(
+            isinstance(item.context_expr, ast.Attribute)
+            and item.context_expr.attr == "plan_stream_ctx"
+            for item in node.items
+        ):
+            continue
+        prior = body[:idx]
+        for prior_node in ast.walk(ast.Module(body=prior, type_ignores=[])):
+            if not isinstance(prior_node, ast.Call):
+                continue
+            callee = prior_node.func
+            if (
+                isinstance(callee, ast.Attribute)
+                and callee.attr == "wait_stream"
+                and isinstance(callee.value, ast.Attribute)
+                and callee.value.attr == "plan_stream"
+            ):
+                return True
+    return False
+
+
+def test_eagle_v2_verify_plan_stream_waits_for_current_stream_inputs():
+    """Plan-stream verify reads tensors produced on the current compute stream."""
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_worker_v2.py")
+    func = _find_function(tree, "verify")
+
+    assert _has_plan_stream_wait_before_context(func)
+
+
+def test_eagle_v2_draft_extend_plan_stream_waits_for_current_stream_inputs():
+    """Draft-extend planning consumes target-verify outputs from current stream."""
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_worker_v2.py")
+    func = _find_function(tree, "_draft_extend_for_decode")
+
+    assert _has_plan_stream_wait_before_context(func)