From a27114d9dcd01260be3946398e638dff9991c64e Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sat, 27 Jun 2026 02:12:37 +0800 Subject: [PATCH] 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. --- .../sglang/srt/speculative/eagle_info_v2.py | 36 +-- .../sglang/srt/speculative/eagle_worker_v2.py | 30 ++- python/sglang/srt/speculative/spec_utils.py | 44 ++++ .../test_eagle_v2_draft_extend_contract.py | 232 ++++++++++++++++++ 4 files changed, 319 insertions(+), 23 deletions(-) create mode 100644 test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 1224fbd33..a095bf988 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 13d4e6c81..0bdfe7209 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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 diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index ed921dafe..5a15a5033 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/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() diff --git a/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py b/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py new file mode 100644 index 000000000..84f5fe2bc --- /dev/null +++ b/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)