[FIX] NSA backend page_table overflow in speculative decoding target_verify (#19016)

2026-03-06 00:04:58 +00:00
parent 13af7cbb02
commit 9ebffef1ef
1 changed files with 3 additions and 1 deletions
--- a/python/sglang/srt/layers/attention/nsa_backend.py
+++ b/python/sglang/srt/layers/attention/nsa_backend.py
@@ -752,9 +752,11 @@ 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
            "page_table": torch.zeros(
                max_num_tokens,
-                self.max_context_len,
+                self.max_context_len + (self.speculative_num_draft_tokens or 0),
                dtype=torch.int32,
                device=self.device,
            ),