chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Prevent spec-v2 decode warmup races

Browse Source 
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
2026-06-27 02:12:37 +08:00
parent 014948c5d7
commit a27114d9dc
4 changed files with 319 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
python/sglang/srt/speculative/eagle_info_v2.py
View File

@@ -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

30
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -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

44
python/sglang/srt/speculative/spec_utils.py
View File

@@ -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()