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

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)

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
+            # Match req_to_token's width exactly. req_to_token is over-allocated
-            # overflow during target_verify when max_seqlen_k = seq_len + num_draft_tokens
+            # 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(
+            # Static width avoids a seq_lens_cpu.max() host read. The captured
-                seq_lens_cpu.max().item() + self.speculative_num_draft_tokens
+            # 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(
+                torch.full(
-                    extend_seq_lens_cpu, dtype=torch.int32, device=self.device
+                    (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)

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 (
+        supports_cuda_draft_extend_graph = False
-            isinstance(self.draft_attn_backend, TritonMultiStepDraftBackend)
+        if _is_cuda:
-            or isinstance(self.draft_attn_backend, TRTLLMMLAMultiStepDraftBackend)
+            # 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 (

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