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

Remove NSA spec-v2 graph metadata host syncs

Browse Source 
NSA spec-v2 draft-extend graph replay was still using host-derived sequence lengths and the draft-decode backend allowlist. That kept NSA on eager draft-extend for spec-v2 and left seq_lens_cpu.max()/list-to-GPU tensor construction on the decode critical path.

This ports the small upstream DSA metadata fixes into the local NSA backend: size the captured graph page table to req_to_token width, use the static captured page-table width for graph replay, split v2 draft-extend from variable-length v1 draft-extend, and decide draft-extend graph support from the prefill-style backend.

Constraint: Current branch does not have the full upstream needs_cpu_seq_lens scheduler/FutureMap infra.

Rejected: Cherry-pick the full DSA fused metadata generation series | too broad and overlaps with local NSA fused metadata-copy code.

Confidence: medium

Scope-risk: moderate

Directive: Do not collapse DRAFT_EXTEND and DRAFT_EXTEND_V2 here; v1 keeps variable accept lengths while v2 must stay graph-static.

Tested: local pytest test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py -q (19 passed)

Tested: local py_compile on nsa_backend.py, nsa_backend_mtp_precompute.py, eagle_worker_v2.py

Tested: remote g0034 cjy-glm5-new pytest test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py -q (19 passed)

Tested: remote g0034 cjy-glm5-new py_compile on nsa_backend.py, nsa_backend_mtp_precompute.py, eagle_worker_v2.py

Not-tested: full decode E2E with SGLANG_ENABLE_SPEC_V2=1
This commit is contained in:
laoyao0822
2026-06-28 02:29:23 +08:00
parent 4414db594c
commit 648a33ab30
4 changed files with 166 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py
View File

@@ -91,9 +91,12 @@ class NativeSparseAttnBackendMTPPrecomputeMixin:
Returns:
PrecomputedMetadata containing all shared intermediate results
"""
# Slice inputs to batch size
# Slice inputs to batch size. seq_lens_cpu may be None when the
# attention backend opts out of the host seq-len mirror; the decode
# replay precompute path no longer needs it.
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
@@ -117,10 +120,13 @@ class NativeSparseAttnBackendMTPPrecomputeMixin:
bs: int,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor],
) -> PrecomputedMetadata:
"""Precompute metadata for normal decode mode."""
max_len = int(seq_lens_cpu.max().item())
# Static page-table width (= captured buffer width) instead of a
# seq_lens_cpu.max() host read. The kernel bounds each row's reads by
# cache_seqlens, so copying the full captured width is correct.
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)

68
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -1186,11 +1186,12 @@ class NativeSparseAttnBackend(
max_bs + 1, dtype=torch.int32, device=self.device
),
# fake page_table for sparse_prefill
# Add extra columns for speculative draft tokens to avoid
# overflow during target_verify when max_seqlen_k = seq_len + num_draft_tokens
# Match req_to_token's width exactly. req_to_token is over-allocated
# beyond max_context_len because spec decode can transiently
# overshoot the nominal context length.
"page_table": torch.zeros(
max_num_tokens,
self.max_context_len + (self.speculative_num_draft_tokens or 0),
self.req_to_token.shape[1],
dtype=torch.int32,
device=self.device,
),
@@ -1375,19 +1376,19 @@ class NativeSparseAttnBackend(
out_cache_loc: Optional[torch.Tensor] = None,
):
"""Initialize forward metadata for replaying CUDA graph."""
assert seq_lens_cpu is not None
self.set_nsa_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: NSAMetadata = self.decode_cuda_graph_metadata[bs]
if forward_mode.is_decode_or_idle():
# Normal Decode
max_len = int(seq_lens_cpu.max().item())
# Static page-table width (= captured buffer width) instead of a
# seq_lens_cpu.max() host read. The kernel bounds each row's reads
# by cache_seqlens, so copying the full captured width is correct.
max_len = metadata.page_table_1.shape[1]
cache_seqlens = seq_lens.to(torch.int32)
metadata.cache_seqlens_int32.copy_(cache_seqlens)
@@ -1402,9 +1403,9 @@ class NativeSparseAttnBackend(
metadata.nsa_cache_seqlens_int32.copy_(nsa_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
)
# Static width avoids a seq_lens_cpu.max() host read. The captured
# page table already uses req_to_token's full allocated width.
max_seqlen_k = metadata.page_table_1.shape[1]
cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to(
torch.int32
@@ -1418,11 +1419,13 @@ class NativeSparseAttnBackend(
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
seqlens_expanded = seqlens_expand_triton(
torch.tensor(
extend_seq_lens_cpu, dtype=torch.int32, device=self.device
torch.full(
(bs,),
self.speculative_num_draft_tokens,
dtype=torch.int32,
device=self.device,
),
cache_seqlens,
self.speculative_num_draft_tokens * bs,
@@ -1433,7 +1436,44 @@ class NativeSparseAttnBackend(
seqlens_expanded, self.nsa_index_topk
)
metadata.nsa_cache_seqlens_int32.copy_(nsa_cache_seqlens)
elif forward_mode.is_draft_extend(include_v2=True):
elif forward_mode.is_draft_extend_v2():
# Spec-v2 draft extend replays the full padded draft width per req.
# Using accept_length here would make graph replay data-dependent
# and reintroduce host syncs; output selection handles accepted
# tokens after verification.
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)
)
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave(
page_indices, repeats=self.speculative_num_draft_tokens, dim=0
)
metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices)
seqlens_expanded = seqlens_expand_triton(
torch.full(
(bs,),
self.speculative_num_draft_tokens,
dtype=torch.int32,
device=self.device,
),
cache_seqlens,
self.speculative_num_draft_tokens * bs,
self.speculative_num_draft_tokens,
)
metadata.nsa_seqlens_expanded.copy_(seqlens_expanded)
nsa_cache_seqlens = compute_nsa_seqlens(
seqlens_expanded, self.nsa_index_topk
)
metadata.nsa_cache_seqlens_int32.copy_(nsa_cache_seqlens)
elif forward_mode.is_draft_extend():
# Spec-v1 draft extend still uses variable accepted lengths.
assert seq_lens_cpu is not None
seq_lens_cpu = seq_lens_cpu[:bs]
max_seqlen_k = int(seq_lens_cpu.max().item())
cache_seqlens = seq_lens.to(torch.int32)
metadata.cache_seqlens_int32.copy_(cache_seqlens)

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

@@ -12,10 +12,6 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_r
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import (
EAGLEDraftNpuGraphRunner,
)
from sglang.srt.layers.attention.triton_backend import TritonMultiStepDraftBackend
from sglang.srt.layers.attention.trtllm_mla_backend import (
TRTLLMMLAMultiStepDraftBackend,
)
from sglang.srt.layers.dp_attention import get_attention_tp_group
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.utils import (
@@ -297,10 +293,28 @@ class EagleDraftWorker(BaseDraftWorker):
self.draft_attn_backend, AiterMultiStepDraftBackend
)
supports_cuda_draft_extend_graph = _is_cuda and (
isinstance(self.draft_attn_backend, TritonMultiStepDraftBackend)
or isinstance(self.draft_attn_backend, TRTLLMMLAMultiStepDraftBackend)
)
supports_cuda_draft_extend_graph = False
if _is_cuda:
# Draft-extend graph compatibility is determined by the prefill-style
# draft_extend backend, not the multi-step draft-decode backend.
# NSA uses NativeSparseAttnBackend here and a separate
# NativeSparseAttnMultiStepBackend for draft decode.
from sglang.srt.layers.attention.nsa_backend import (
NativeSparseAttnBackend,
)
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
from sglang.srt.layers.attention.trtllm_mla_backend import (
TRTLLMMLABackend,
)
supports_cuda_draft_extend_graph = isinstance(
self.draft_extend_attn_backend,
(
TritonAttnBackend,
TRTLLMMLABackend,
NativeSparseAttnBackend,
),
)
# Capture extend
# TODO: support draft extend cuda graph for more attention backends
if self.draft_extend_attn_backend and (

80
test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py
View File

@@ -163,6 +163,86 @@ def test_eagle_v2_binds_draft_runner_to_draft_extend_attention_backend():
raise AssertionError("draft_runner.attn_backend is not bound to draft_extend backend")
def test_nsa_cuda_graph_page_table_is_sized_to_req_to_token_width():
"""NSA graph page tables must match req_to_token's allocated width.
Spec decode can transiently exceed max_context_len. Sizing the captured
page table from max_context_len makes replay copy widths depend on a
narrower fake table than req_to_token actually owns.
"""
tree = _parse_module("python/sglang/srt/layers/attention/nsa_backend.py")
cls = _find_class(tree, "NativeSparseAttnBackend")
func = _find_method(cls, "init_cuda_graph_state")
text = ast.unparse(func)
assert "self.req_to_token.shape[1]" in text
assert (
"self.max_context_len + (self.speculative_num_draft_tokens or 0)"
not in text
)
def test_nsa_cuda_graph_replay_uses_static_page_table_width():
"""Replay metadata should not read seq_lens_cpu to choose copy width.
CUDA graph replay can copy the captured page-table width because the NSA
kernels bound row access with cache_seqlens. Reading seq_lens_cpu.max()
introduces a host sync in the decode hot path.
"""
tree = _parse_module("python/sglang/srt/layers/attention/nsa_backend.py")
cls = _find_class(tree, "NativeSparseAttnBackend")
func = _find_method(cls, "init_forward_metadata_replay_cuda_graph")
text = ast.unparse(func)
graph_sync_free_prefix = text.split("elif forward_mode.is_draft_extend():", 1)[0]
assert "seq_lens_cpu.max().item()" not in graph_sync_free_prefix
assert "metadata.page_table_1.shape[1]" in text
assert "assert seq_lens_cpu is not None" not in graph_sync_free_prefix
def test_nsa_cuda_graph_replay_keeps_draft_extend_v1_v2_semantics_separate():
"""Spec-v2 draft-extend graph uses static width; spec-v1 stays variable.
The v2 graph path must not use accept_length/tolist to resize replay
tensors. The v1 path still needs variable accepted lengths.
"""
tree = _parse_module("python/sglang/srt/layers/attention/nsa_backend.py")
cls = _find_class(tree, "NativeSparseAttnBackend")
func = _find_method(cls, "init_forward_metadata_replay_cuda_graph")
text = ast.unparse(func)
assert "forward_mode.is_draft_extend_v2()" in text
assert "forward_mode.is_draft_extend(include_v2=True)" not in text
assert "torch.full((bs,), self.speculative_num_draft_tokens" in text
assert "extend_seq_lens = spec_info.accept_length[:bs]" in text
assert "elif forward_mode.is_draft_extend():" in text
v1_branch = text.split("elif forward_mode.is_draft_extend():", 1)[1]
assert "seq_lens_cpu.max().item()" in v1_branch
def test_eagle_v2_cuda_draft_extend_graph_allows_nsa_prefill_backend():
"""Spec-v2 should decide draft-extend graph support from prefill backend.
NSA creates a multi-step backend for draft decode and a NativeSparseAttnBackend
for draft extend. Checking the decode backend keeps NSA on the eager
draft-extend path even after the replay metadata is graph-safe.
"""
tree = _parse_module("python/sglang/srt/speculative/eagle_worker_v2.py")
cls = _find_class(tree, "EagleDraftWorker")
func = _find_method(cls, "init_cuda_graphs")
text = ast.unparse(func)
assert "NativeSparseAttnBackend" in text
assert "self.draft_extend_attn_backend" in text
assert "self.draft_attn_backend, TritonMultiStepDraftBackend" not in text
assert "self.draft_attn_backend, TRTLLMMLAMultiStepDraftBackend" not in text
def _function_calls(func: ast.FunctionDef, name: str) -> int:
count = 0