From 648a33ab30d6fd65ead25b1f4d633c19eab677f2 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sun, 28 Jun 2026 02:29:23 +0800 Subject: [PATCH] Remove NSA spec-v2 graph metadata host syncs 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 --- .../nsa/nsa_backend_mtp_precompute.py | 14 +++- .../srt/layers/attention/nsa_backend.py | 68 ++++++++++++---- .../sglang/srt/speculative/eagle_worker_v2.py | 30 +++++-- .../test_eagle_v2_draft_extend_contract.py | 80 +++++++++++++++++++ 4 files changed, 166 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py index 846d276e4..e84a2b9cc 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_backend_mtp_precompute.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 4faf58c8f..43019ab6e 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index f68ec3b8d..c3007d3a0 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -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 ( diff --git a/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py b/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py index 70e5a9a95..b5b4b4bbe 100644 --- a/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py +++ b/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py @@ -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