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

[DSA] Enable draft-extend CUDA graph for DeepSeek Sparse Attention (#29413)

Browse Source 
Co-authored-by: kpham-sgl <kpham-sgl@users.noreply.github.com>
Co-authored-by: Jackey Hua <107608053+zhendonghua@users.noreply.github.com>
This commit is contained in:
Khoa Pham
2026-06-26 23:53:46 -07:00
committed by GitHub
co-authored by kpham-sgl Jackey Hua
parent 81d3c3ce77
commit 9214b9338f
3 changed files with 70 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py
+3 -2
View File
@@ -90,7 +90,8 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
"""
# Slice inputs to batch size
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
if seq_lens_cpu is not None:
seq_lens_cpu = seq_lens_cpu[:bs]
req_pool_indices = req_pool_indices[:bs]
# Dispatch to mode-specific precomputation
@@ -113,7 +114,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
seq_lens_cpu: torch.Tensor,
) -> PrecomputedMetadata:
"""Precompute metadata for normal decode mode."""
max_len = int(seq_lens_cpu.max().item())
max_len = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
# Convert to int32 and compute cumsum
cache_seqlens = seq_lens.to(torch.int32)
python/sglang/srt/layers/attention/dsa_backend.py
+50 -28
View File
@@ -303,6 +303,11 @@ _DSA_IMPL_T: TypeAlias = Literal[
class DeepseekSparseAttnBackend(
DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend
):
# Decode/verify/draft graph replay rebuilds metadata from static buffers
# (page-table width) and never reads seq_lens_cpu / seq_lens_sum; opt out of
# the D2H sync. The eager fallback derives lengths from GPU seq_lens.
needs_cpu_seq_lens: bool = False
def __init__(
self,
model_runner: ModelRunner,
@@ -642,8 +647,15 @@ class DeepseekSparseAttnBackend(
cache_seqlens_int32 = (forward_batch.seq_lens + draft_token_num).to(torch.int32)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32)
assert forward_batch.seq_lens_cpu is not None
max_seqlen_k = int(forward_batch.seq_lens_cpu.max().item() + draft_token_num)
if forward_batch.seq_lens_cpu is not None:
max_seqlen_k = int(
forward_batch.seq_lens_cpu.max().item() + draft_token_num
)
else:
# needs_cpu_seq_lens=False nulls the host mirror for spec-v2 relay
# batches; graph replay uses the static page-table width, so only this
# eager (e.g. over-capture-bs) fallback needs a length here.
max_seqlen_k = int(forward_batch.seq_lens.max().item()) + draft_token_num
# [b, max_seqlen_k]
page_table = self.req_to_token_pool.req_to_token[
forward_batch.req_pool_indices, :max_seqlen_k
@@ -1187,8 +1199,6 @@ class DeepseekSparseAttnBackend(
also call this directly via _apply_cuda_graph_metadata when they
need to pass out_cache_loc / actual_forward_mode explicitly.
"""
assert seq_lens_cpu is not None
if bs not in self.decode_cuda_graph_metadata:
self._build_forward_metadata_cuda_graph(
bs,
@@ -1206,14 +1216,13 @@ class DeepseekSparseAttnBackend(
self.set_dsa_prefill_impl(forward_batch=None)
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[:bs]
req_pool_indices = req_pool_indices[:bs]
# Normal Decode
metadata: DSAMetadata = self.decode_cuda_graph_metadata[bs]
if forward_mode.is_decode_or_idle():
# Normal Decode
max_len = int(seq_lens_cpu.max().item())
max_len = metadata.page_table_1.shape[1]
cache_seqlens = seq_lens.to(torch.int32)
metadata.cache_seqlens_int32.copy_(cache_seqlens)
@@ -1228,9 +1237,7 @@ class DeepseekSparseAttnBackend(
metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens)
seqlens_expanded = cache_seqlens
elif forward_mode.is_target_verify():
max_seqlen_k = int(
seq_lens_cpu.max().item() + self.speculative_num_draft_tokens
)
max_seqlen_k = metadata.page_table_1.shape[1]
cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to(
torch.int32
@@ -1244,12 +1251,18 @@ class DeepseekSparseAttnBackend(
page_indices, repeats=self.speculative_num_draft_tokens, dim=0
)
metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices)
extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs
# Fill the constant per-req qo lengths (num_draft_tokens) on-device;
# torch.tensor(list, device=cuda) does a pageable H2D copy that
# blocks the host on the whole queued stream.
extend_seq_lens = torch.full(
(bs,),
self.speculative_num_draft_tokens,
dtype=torch.int32,
device=self.device,
)
seqlens_expanded = seqlens_expand_triton(
torch.tensor(
extend_seq_lens_cpu, dtype=torch.int32, device=self.device
),
extend_seq_lens,
cache_seqlens,
self.speculative_num_draft_tokens * bs,
self.speculative_num_draft_tokens,
@@ -1260,39 +1273,44 @@ class DeepseekSparseAttnBackend(
)
metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens)
elif forward_mode.is_draft_extend_v2():
max_seqlen_k = int(seq_lens_cpu.max().item())
# V2 draft-extend processes the full padded tree width
# (speculative_num_draft_tokens) per req -- a static shape, like
# target-verify -- so graph replay stays host-sync-free. seq_lens
# already includes the draft KV written by prepare_for_draft_extend;
# the per-req accept length is handled downstream by output
# selection, not by reshaping the page table here.
max_seqlen_k = metadata.page_table_1.shape[1]
cache_seqlens = seq_lens.to(torch.int32)
metadata.cache_seqlens_int32.copy_(cache_seqlens)
metadata.cu_seqlens_k[1:].copy_(
torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32)
)
extend_seq_lens = spec_info.num_accept_tokens[:bs]
extend_seq_lens_cpu = extend_seq_lens.tolist()
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave(
page_indices, repeats=extend_seq_lens, dim=0
)
metadata.page_table_1[: page_indices.shape[0], :max_seqlen_k].copy_(
page_indices
page_indices, repeats=self.speculative_num_draft_tokens, dim=0
)
metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices)
# See target-verify note: fill on-device to avoid the blocking
# pageable H2D from torch.tensor(list, device=cuda).
extend_seq_lens = torch.full(
(bs,),
self.speculative_num_draft_tokens,
dtype=torch.int32,
device=self.device,
)
seqlens_expanded = seqlens_expand_triton(
extend_seq_lens,
cache_seqlens,
sum(extend_seq_lens_cpu),
self.speculative_num_draft_tokens * bs,
self.speculative_num_draft_tokens,
)
metadata.dsa_seqlens_expanded[: seqlens_expanded.shape[0]].copy_(
seqlens_expanded
)
metadata.dsa_seqlens_expanded.copy_(seqlens_expanded)
dsa_cache_seqlens = compute_dsa_seqlens(
seqlens_expanded, self.dsa_index_topk
)
metadata.dsa_cache_seqlens_int32[: seqlens_expanded.shape[0]].copy_(
dsa_cache_seqlens
)
metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens)
# Update DeepGEMM paged MQA schedule metadata outside the captured graph.
if is_cuda() and (
@@ -2588,6 +2606,10 @@ class DeepseekSparseAttnBackend(
class DeepseekSparseAttnMultiStepBackend:
# Per-step draft decode replays from precomputed GPU metadata; opt out so
# decide_needs_cpu_seq_lens' OR over the backends stays False.
needs_cpu_seq_lens: bool = False
def __init__(
self, model_runner: ModelRunner, topk: int, speculative_num_steps: int
):
python/sglang/srt/speculative/eagle_worker_v2.py
+17 -7
View File
@@ -424,15 +424,25 @@ class EagleDraftWorker(EagleDraftWorkerBase):
self.draft_attn_backend, AiterMultiStepDraftBackend
)
graph_supported_backend_types = [
TritonAttnBackend,
TRTLLMMLABackend,
TRTLLMHAAttnBackend,
TokenspeedMLABackend,
FlashInferAttnBackend,
]
if _is_cuda or _is_musa:
# DSA is CUDA-only; import lazily so non-CUDA builds don't pull in
# deep_gemm and the rest of the sparse-attention stack at import time.
from sglang.srt.layers.attention.dsa_backend import (
DeepseekSparseAttnBackend,
)
graph_supported_backend_types.append(DeepseekSparseAttnBackend)
graph_supported_backend = isinstance(
self.draft_extend_attn_backend,
(
TritonAttnBackend,
TRTLLMMLABackend,
TRTLLMHAAttnBackend,
TokenspeedMLABackend,
FlashInferAttnBackend,
),
tuple(graph_supported_backend_types),
)
supports_cuda_draft_extend_graph = (
_is_cuda or _is_musa