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.
This commit is contained in:
laoyao0822
@@ -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)
|
||||
Reference in New Issue
Block a user