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

[DSV4] Remove per-step seqlen D2H from speculative to make overlap scheduler work (#30365)

Browse Source 
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-07-14 18:01:22 -07:00
committed by GitHub
co-authored by weireweire
parent b8a00e2ec8
commit a9cf5e68e6
5 changed files with 115 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+64 -48
View File
@@ -68,7 +68,7 @@ from sglang.srt.speculative.ragged_verify import (
read_ragged_verify_mode,
resolve_ragged_verify_layout,
)
from sglang.srt.utils import ceil_align, is_xpu
from sglang.srt.utils import ceil_align, is_cuda, is_xpu
from sglang.srt.utils.common import is_sm120_supported
if TYPE_CHECKING:
@@ -79,6 +79,7 @@ if TYPE_CHECKING:
from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
_is_sm120 = is_sm120_supported()
_is_cuda = is_cuda()
_is_xpu = is_xpu()
logger = logging.getLogger(__name__)
@@ -497,6 +498,7 @@ class DeepseekV4AttnBackend(
super().__init__()
self.model_runner = model_runner
self.device = torch.device(model_runner.device)
self.max_context_len = model_runner.model_config.context_len
head_dim = model_runner.model_config.head_dim
assert (
head_dim == 512
@@ -553,15 +555,13 @@ class DeepseekV4AttnBackend(
DSV4RawDecodeMetadata,
] = None
self.online_c128_mtp = OnlineC128MTPController(self)
# Draft-extend and online-c128 verify metadata are host-planned, so
# spec runs keep the relay publish (the mirror only exists under
# spec-v2; without spec the flag has no consumer either way).
# DSPARK is the exception: its draft path carries its own host lens
# (reserved_seq_lens_cpu) and its verify prep is device-side.
spec_alg = model_runner.spec_algorithm
if not spec_alg.is_none() and not spec_alg.is_dspark():
self.needs_cpu_seq_lens = True
self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device)
spec_alg = model_runner.spec_algorithm
self.needs_cpu_seq_lens = not spec_alg.is_dspark() and (
not _is_cuda
or not envs.SGLANG_PREP_IN_CUDA_GRAPH.get()
or self.online_c128_mtp.enabled()
)
self.is_dspark_draft = model_runner.is_draft_worker and spec_alg.is_dspark()
@@ -610,7 +610,7 @@ class DeepseekV4AttnBackend(
self,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: List[int],
seq_lens_cpu: Optional[List[int]],
extend_seq_lens: torch.Tensor,
use_prefill_cuda_graph: bool,
online_c128_state_slot_offset: int,
@@ -618,6 +618,7 @@ class DeepseekV4AttnBackend(
if not self.online_c128_mtp.enabled():
return None
assert seq_lens_cpu is not None
num_draft_tokens = self.speculative_num_draft_tokens
seq_lens_cpu = [int(x) + num_draft_tokens for x in seq_lens_cpu]
extend_lens_cpu = [num_draft_tokens] * len(seq_lens_cpu)
@@ -800,9 +801,11 @@ class DeepseekV4AttnBackend(
if envs.SGLANG_PREP_IN_CUDA_GRAPH.get():
assert out_cache_loc is not None
bs = len(seq_lens)
seq_lens_cpu_list = (
seq_lens_cpu.tolist() if seq_lens_cpu is not None else None
)
if self.needs_cpu_seq_lens:
assert seq_lens_cpu is not None
seq_lens_cpu_list = seq_lens_cpu.tolist()
else:
seq_lens_cpu_list = None
if ragged_layout is None:
self.extend_seq_lens_buffer[:bs].fill_(
self.speculative_num_draft_tokens
@@ -838,12 +841,14 @@ class DeepseekV4AttnBackend(
total_verify_tokens=total_verify_tokens,
)
else:
seq_lens_cpu = seq_lens.tolist()
seq_lens_cpu_list = (
seq_lens_cpu.tolist() if seq_lens_cpu is not None else seq_lens.tolist()
)
return self.init_forward_metadata_target_verify_old(
max_seq_len=max_seq_len,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_cpu=seq_lens_cpu,
seq_lens_cpu=seq_lens_cpu_list,
out_cache_loc=out_cache_loc,
use_prefill_cuda_graph=use_prefill_cuda_graph,
online_c128_state_slot_offset=online_c128_state_slot_offset,
@@ -1039,32 +1044,38 @@ class DeepseekV4AttnBackend(
def init_forward_metadata_draft_extend(
self,
max_seq_len: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: List[int],
num_tokens_per_req: int,
out_cache_loc: Optional[torch.Tensor] = None,
use_prefill_cuda_graph: bool = False,
) -> DSV4Metadata:
batch_size = len(seq_lens)
extend_seq_lens_cpu = [num_tokens_per_req] * batch_size
extend_seq_lens = self._move_to_device(extend_seq_lens_cpu)
num_tokens = num_tokens_per_req * batch_size
if out_cache_loc is None:
out_cache_loc = seq_lens.new_zeros(num_tokens)
return self.init_forward_metadata_prefill(
seq_lens=seq_lens,
max_seq_len=max_seq_len,
req_pool_indices=req_pool_indices,
seq_lens_cpu=seq_lens_cpu,
out_cache_loc=out_cache_loc,
num_tokens=num_tokens,
extend_seq_lens=extend_seq_lens,
extend_seq_lens_cpu=extend_seq_lens_cpu,
extend_start_loc=None,
seq_lens_casual, req_pool_indices_repeated = (
self.expand_extend_with_same_length(
bs=batch_size,
qo_len=num_tokens_per_req,
seq_lens=seq_lens,
req_pool_indices=req_pool_indices,
)
)
core_attn_metadata = self.make_core_attn_metadata(
req_to_token=self.req_to_token,
req_pool_indices_repeated=req_pool_indices_repeated,
seq_lens_casual=seq_lens_casual,
# Draft extend is SWA-only. Keep the required 2-D page-table
# placeholder narrow instead of materializing the full context.
max_seq_len=self.page_size,
out_loc=out_cache_loc,
need_compress=False,
use_prefill_cuda_graph=use_prefill_cuda_graph,
is_prefill=True,
)
return DSV4Metadata(
core_attn_metadata=core_attn_metadata,
indexer_metadata=None,
)
def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None:
@@ -1149,6 +1160,7 @@ class DeepseekV4AttnBackend(
bs = forward_batch.batch_size
req_pool_indices = forward_batch.req_pool_indices
seq_lens = forward_batch.seq_lens
uses_cpu_seq_lens = self.needs_cpu_seq_lens or self.is_dspark_draft
if in_capture:
# Captured graph does no real cache writes, so synthesize a dummy
@@ -1164,14 +1176,14 @@ class DeepseekV4AttnBackend(
out_cache_loc = None
actual_forward_mode = forward_batch.forward_mode
seq_lens_sum = int(seq_lens.sum().item())
seq_lens_cpu = seq_lens.cpu()
seq_lens_cpu = seq_lens.cpu() if uses_cpu_seq_lens else None
else:
out_cache_loc = forward_batch.out_cache_loc
actual_forward_mode = getattr(
forward_batch, "actual_forward_mode", forward_batch.forward_mode
)
seq_lens_sum = forward_batch.seq_lens_sum
seq_lens_cpu = forward_batch.seq_lens_cpu
seq_lens_cpu = forward_batch.seq_lens_cpu if uses_cpu_seq_lens else None
if actual_forward_mode == ForwardMode.IDLE:
logger.debug(
@@ -1181,7 +1193,8 @@ class DeepseekV4AttnBackend(
)
device = seq_lens.device
seq_lens = torch.ones(bs, dtype=seq_lens.dtype, device=device)
seq_lens_cpu = torch.ones(bs, dtype=torch.int64)
if uses_cpu_seq_lens:
seq_lens_cpu = torch.ones(bs, dtype=torch.int64)
seq_lens_sum = bs
req_pool_indices = torch.zeros(
bs, dtype=req_pool_indices.dtype, device=device
@@ -1297,17 +1310,11 @@ class DeepseekV4AttnBackend(
mode="constant",
value=0,
)
draft_extend_seq_lens_cpu = (
seq_lens_cpu.tolist() if seq_lens_cpu is not None else seq_lens.tolist()
)
temp_metadata = self.init_forward_metadata_draft_extend(
max_seq_len=chosen_max_seq_len,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_cpu=draft_extend_seq_lens_cpu,
num_tokens_per_req=num_tokens_per_req,
out_cache_loc=out_cache_loc,
use_prefill_cuda_graph=True,
)
else:
self.online_c128_mtp.clear()
@@ -1348,6 +1355,9 @@ class DeepseekV4AttnBackend(
logical_forward_mode = _get_logical_forward_mode(forward_batch)
req_pool_indices = forward_batch.req_pool_indices
seq_lens = forward_batch.seq_lens.to(torch.int32)
# Regular prefill batches already carry scheduler-maintained CPU lengths.
# Keep using those when present; needs_cpu_seq_lens only controls whether
# speculative overlap must publish a new GPU-to-CPU mirror each step.
seq_lens_cpu = forward_batch.seq_lens_cpu
assert self.req_to_token_pool.req_to_token is self.req_to_token
@@ -1359,7 +1369,7 @@ class DeepseekV4AttnBackend(
elif seq_lens_cpu is not None:
max_seq_len = int(seq_lens_cpu.max().item())
else:
max_seq_len = int(seq_lens.max().item())
max_seq_len = self.MAX_SEQ_LEN_FOR_CAPTURE
verify_bs = _get_target_verify_bs(forward_batch)
online_c128_state_slot_offset = self.online_c128_mtp.prepare_forward(
logical_forward_mode,
@@ -1406,29 +1416,35 @@ class DeepseekV4AttnBackend(
online_c128_state_slot_offset=online_c128_state_slot_offset,
ragged_layout=ragged_layout,
)
elif logical_forward_mode.is_prefill(include_draft_extend_v2=True):
elif logical_forward_mode.is_draft_extend_v2():
num_tokens_per_req = self.speculative_num_draft_tokens
assert num_tokens_per_req > 0
metadata = self.init_forward_metadata_draft_extend(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
num_tokens_per_req=num_tokens_per_req,
out_cache_loc=forward_batch.out_cache_loc,
)
elif logical_forward_mode.is_prefill():
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
extend_seq_lens = forward_batch.extend_seq_lens
assert (
seq_lens is not None
and seq_lens_cpu is not None
and extend_seq_lens is not None
and extend_seq_lens_cpu is not None
)
is_draft = forward_batch.forward_mode.is_draft_extend_v2()
prefill_seq_lens_cpu = (
seq_lens_cpu.tolist() if seq_lens_cpu is not None else seq_lens.tolist()
)
metadata = self.init_forward_metadata_prefill(
max_seq_len=max_seq_len,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_cpu=prefill_seq_lens_cpu,
seq_lens_cpu=seq_lens_cpu.tolist(),
out_cache_loc=forward_batch.out_cache_loc,
num_tokens=sum(extend_seq_lens_cpu),
extend_seq_lens=extend_seq_lens,
extend_seq_lens_cpu=extend_seq_lens_cpu,
extend_start_loc=forward_batch.extend_start_loc,
need_compress=not is_draft,
need_compress=True,
use_prefill_cuda_graph=use_prefill_cuda_graph,
)
else:
python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py
+6
View File
@@ -1490,6 +1490,7 @@ def run_dsv4_target_verify_attention_case(
fixture = build_dsv4_attention_fixture(testcase, case, dtype=dtype, device=device)
runner = fixture.runner
max_context_len = runner.req_to_token_pool.req_to_token.shape[1]
testcase.assertEqual(fixture.backend.max_context_len, max_context_len)
_populate_swa_kv_cache(fixture, max_context_len=max_context_len, device=device)
if case.compress_ratio in (4, 128):
@@ -1530,6 +1531,7 @@ def run_dsv4_draft_extend_attention_case(
*,
dtype: torch.dtype = torch.bfloat16,
device: str = "cuda",
force_gpu_only_seq_lens: bool = False,
) -> None:
"""Math-faithful EAGLE `DRAFT_EXTEND` test for DSV4.
@@ -1567,6 +1569,10 @@ def run_dsv4_draft_extend_attention_case(
fixture.forward_batch,
device=device,
)
if force_gpu_only_seq_lens:
fixture.forward_batch.seq_lens_cpu = None
fixture.forward_batch.seq_lens_sum = None
fixture.forward_batch.spec_info.seq_lens_cpu = None
q_input, _ = fixture.actual_module.project(fixture.input_hidden)
with torch.no_grad(), forward_context(ForwardContext(attn_backend=fixture.backend)):
python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py
+10 -1
View File
@@ -1484,6 +1484,7 @@ def run_dsv4_eagle_draft_cuda_graph_runner_case(
vocab_size: int = 64,
dtype: torch.dtype = torch.bfloat16,
device: str = "cuda",
force_gpu_only_seq_lens: bool = False,
):
settings = EagleDraftRunnerSettings(
topk=topk,
@@ -1498,12 +1499,20 @@ def run_dsv4_eagle_draft_cuda_graph_runner_case(
atol=DSV4_ATOL,
rtol=DSV4_RTOL,
)
def _make_forward_batch(case, draft_inputs, settings):
batch = _make_dsv4_eagle_draft_forward_batch(case, draft_inputs, settings)
if force_gpu_only_seq_lens:
batch.seq_lens_cpu = None
batch.seq_lens_sum = None
return batch
adapter = EagleDraftCudaGraphRunnerAdapter(
build_fixture=build_dsv4_attention_fixture,
make_model_forward=_make_dsv4_model_forward,
make_draft_inputs=_make_dsv4_draft_inputs,
prepare_replay_state=_prepare_dsv4_draft_replay_state,
make_forward_batch=_make_dsv4_eagle_draft_forward_batch,
make_forward_batch=_make_forward_batch,
check_case=_check_dsv4_draft_cache_layout,
init_eager_metadata=_init_dsv4_eager_metadata,
)
python/sglang/test/kits/attention_unittest/runner_modes/speculative_target_verify_runner.py
+6
View File
@@ -890,6 +890,7 @@ def run_dsv4_eagle_verify_cuda_graph_case(
dtype: torch.dtype = torch.bfloat16,
device: str = "cuda",
cuda_graph_capture_batch_size: int = 2,
force_gpu_only_seq_lens: bool = False,
):
"""DSV4 EAGLE target_verify CUDA-graph capture/replay. Chain only —
`DeepseekV4AttnBackend.__init__` asserts `self.topk in [0, 1]` at
@@ -936,6 +937,11 @@ def run_dsv4_eagle_verify_cuda_graph_case(
batch.spec_info = _make_eagle_verify_input(
spec_case, batch, topk=topk, device=device
)
if force_gpu_only_seq_lens:
batch.seq_lens_cpu = None
batch.seq_lens_sum = None
batch.spec_info.seq_lens_cpu = None
batch.spec_info.seq_lens_sum = None
def _make_capture_case(base, name, capture_prefix_len: int, bs: int):
# Capture uses uniform prefixes per request; each request still