Skip to content

Fix DSv4 attention backend for EAGLE per-step draft #24750

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ch-wan wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix DSv4 attention backend for EAGLE per-step draft #24750

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion python/sglang/srt/layers/attention/deepseek_v4_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -664,11 +664,20 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None:
max_seq_len = int(seq_lens_cpu.max().item())

if forward_batch.forward_mode.is_decode_or_idle():
out_cache_loc = forward_batch.out_cache_loc
if self.speculative_num_steps > 0:
Copy link
Copy Markdown
Collaborator

@Fridge003 Fridge003 May 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe only do this when cuda graph is disabled

# Per-step EAGLE draft: slice the shared `out_cache_loc` to
# this step's portion and advance `seq_lens` by `step_id + 1`.
bs = req_pool_indices.shape[0]
step_id = self.speculative_step_id
out_cache_loc = out_cache_loc[bs * step_id : bs * (step_id + 1)]
seq_lens = seq_lens + (step_id + 1)
max_seq_len = max_seq_len + (step_id + 1)
metadata = self.init_forward_metadata_decode(
max_seq_len=max_seq_len,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
out_cache_loc=forward_batch.out_cache_loc,
out_cache_loc=out_cache_loc,
)
elif forward_batch.forward_mode.is_target_verify():
metadata = self.init_forward_metadata_target_verify(
Expand Down Expand Up @@ -813,6 +822,14 @@ def init_forward_metadata_replay_cuda_graph(
if bucket == _GraphBucket.DECODE_OR_IDLE:
assert out_cache_loc is not None
assert len(out_cache_loc.shape) == 1, f"{out_cache_loc.shape=}"
if self.speculative_num_steps > 0:
Copy link
Copy Markdown
Collaborator

@Fridge003 Fridge003 May 11, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the bug only happens when disabling cuda graph, these lines are unnecessary?

# Per-step EAGLE draft: take this step's slice of the
# shared `out_cache_loc` and advance `seq_lens` by
# `step_id + 1`. The captured buffer's `chosen_max_seq_len`
# already accommodates the spec extension.
step_id = self.speculative_step_id
out_cache_loc = out_cache_loc[bs * step_id : bs * (step_id + 1)]
seq_lens = seq_lens + (step_id + 1)
out_cache_loc_padded = torch.nn.functional.pad(
out_cache_loc,
pad=(0, bs - len(out_cache_loc)),
Expand Down
Loading