[Spec] Compact the target-verify mask when nothing reads it (#32920)

Co-authored-by: Kaixi Matteo Chen <kaiximatteoc@nvidia.com>

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -397,14 +397,6 @@ class AscendAttnBackend(AttentionBackend):
         v = layer.v_head_dim
         return (d == v and d in (128, 192, 256)) or (d == 192 and v == 128)
 
-    def get_verify_buffers_to_fill_after_draft(self):
-        """
-        Return buffers for verify attention kernels that needs to be filled after draft.
-
-        Typically, these are tree mask and position buffers.
-        """
-        return [None, None]
-
     def update_verify_buffers_to_fill_after_draft(
         self, spec_info: SpecInput, cuda_graph_bs: Optional[int]
     ):

python/sglang/srt/layers/attention/base_attn_backend.py

@@ -10,6 +10,7 @@ from sglang.srt.utils.common import is_npu
 
 if TYPE_CHECKING:
     from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata
+    from sglang.srt.layers.attention.verify_mask import VerifyMask
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
     from sglang.srt.speculative.spec_info import SpecInput
@@ -166,21 +167,10 @@ class AttentionBackend(ABC):
         """
         pass
 
-    def get_verify_buffers_to_fill_after_draft(self):
+    @property
-        """
+    def verify_mask(self) -> Optional[VerifyMask]:
-        Return buffers of verify attention kernels that needs to be filled after draft.
+        """The mask the draft stage fills in place, if this backend has one."""
-
+        return None
-        Typically, these are tree mask and position buffers.
-        """
-        return [None, None]
-
-    def target_verify_reads_custom_mask(self) -> bool:
-        """Whether target-verify attention reads spec_info.custom_mask at all.
-
-        When False, build_tree_kernel_efficient skips the full-buffer prefix
-        fill (max_num_tokens x max_context_len bool memset per verify step).
-        """
-        return True
 
     def update_verify_buffers_to_fill_after_draft(
         self, spec_info: SpecInput, cuda_graph_bs: Optional[int]

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -57,6 +57,7 @@ from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
     SparsePrefillChunkCache,
     SparsePrefillWorkspace,
 )
+from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_parallel
@@ -571,7 +572,7 @@ class DeepseekV4AttnBackend(
 
         self.is_dspark_draft = model_runner.is_draft_worker and spec_alg.is_dspark()
         self.is_draft_runner = model_runner.is_draft_worker
-        self.cuda_graph_custom_mask = None
+        self._verify_mask = None
 
     def _move_to_device(self, x: List[int]) -> torch.Tensor:
         pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True)
@@ -1509,24 +1510,20 @@ class DeepseekV4AttnBackend(
         self.draft_extend_num_tokens_per_req = (
             max_num_tokens // max_bs if max_bs > 0 else 1
         )
-        if self.speculative_num_draft_tokens and not self.is_draft_runner:
+        # Verify metadata never extracts the mask. No skip_prefill notion here.
-            # DSV4's verify metadata ignores custom_mask, but handing
+        self._verify_mask = maybe_create_verify_mask(
-            # build_tree a preallocated scratch keeps it from dynamically
+            is_draft_runner=self.is_draft_runner,
-            # allocating a FULL_MASK buffer (bs * max_context_len under the
+            skip_prefill=False,
-            # GPU-only spec path) every verify step.
+            max_bs=max_bs,
-            self.cuda_graph_custom_mask = torch.zeros(
+            max_context_len=self.max_context_len,
-                max_num_tokens
+            num_draft_tokens=self.speculative_num_draft_tokens,
-                * (self.max_context_len + self.speculative_num_draft_tokens),
-                dtype=torch.bool,
             device=self.device,
+            is_read=False,
         )
 
-    def get_verify_buffers_to_fill_after_draft(self):
+    @property
-        return [self.cuda_graph_custom_mask, None]
+    def verify_mask(self) -> Optional[VerifyMask]:
-
+        return self._verify_mask
-    def target_verify_reads_custom_mask(self) -> bool:
-        # DSV4 verify metadata never extracts from custom_mask.
-        return False
 
     def replay_cuda_graph_metadata_from(
         self,

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -18,6 +18,7 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
 )
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
 from sglang.srt.layers.cp.utils import is_cp_v2_active
 from sglang.srt.layers.radix_attention import AttentionType
@@ -179,15 +180,13 @@ class FlashAttentionBackend(AttentionBackend):
             self.max_context_len + self.page_size - 1
         ) // self.page_size
         # Page table is built on-device (build_trtllm_mha_page_table) and the
-        # tree-mask scratch is preallocated (get_verify_buffers_to_fill_after_draft),
+        # tree mask is preallocated (see VerifyMask), so no
-        # so no seq_lens_cpu / seq_lens_sum D2H sync is ever needed.
+        # seq_lens_cpu / seq_lens_sum D2H sync is ever needed.
         self.needs_cpu_seq_lens = False
         self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
         self.skip_prefill = skip_prefill
         self.attn_cp_size = model_runner.ps.attn_cp_size
-        # Preallocated FULL_MASK tree-mask scratch; lets build_tree_kernel_efficient
+        self._verify_mask = None
-        # avoid the seq_lens_sum D2H sync (see get_verify_buffers_to_fill_after_draft).
-        self.cuda_graph_custom_mask = None
         # The worker fetches the tree-mask scratch from the target backend
         # only; draft-side instances must not allocate it.
         self.is_draft_runner = model_runner.is_draft_worker
@@ -2199,16 +2198,15 @@ class FlashAttentionBackend(AttentionBackend):
                 ),
             }
 
-            # Worst-case FULL_MASK tree-mask scratch (bool). build_tree_kernel
+            # topk<=1 never extracts the mask; both metadata paths gate on topk > 1.
-            # fills it in-place, so the GPU-only path needs no seq_lens_sum.
+            self._verify_mask = maybe_create_verify_mask(
-            # Costs max_num_tokens * max_context_len bytes (can reach 100s of
+                is_draft_runner=self.is_draft_runner,
-            # MB at long context) and is fully memset every verify step.
+                skip_prefill=self.skip_prefill,
-            if not self.skip_prefill and not self.is_draft_runner:
+                max_bs=max_bs,
-                self.cuda_graph_custom_mask = torch.zeros(
+                max_context_len=self.max_context_len,
-                    max_num_tokens
+                num_draft_tokens=self.speculative_num_draft_tokens,
-                    * (self.max_context_len + self.speculative_num_draft_tokens),
-                    dtype=torch.bool,
                 device=self.device,
+                is_read=self.topk > 1,
             )
 
             self.draft_extend_metadata = {
@@ -2557,16 +2555,9 @@ class FlashAttentionBackend(AttentionBackend):
 
         return metadata, metadata_expand
 
-    def get_verify_buffers_to_fill_after_draft(self):
+    @property
-        # Return the preallocated FULL_MASK tree-mask scratch so that
+    def verify_mask(self) -> Optional[VerifyMask]:
-        # build_tree_kernel_efficient fills it in-place and the worker never
+        return self._verify_mask
-        # needs seq_lens_sum to size a dynamic allocation (no D2H sync).
-        return [self.cuda_graph_custom_mask, None]
-
-    def target_verify_reads_custom_mask(self) -> bool:
-        # topk<=1 verify never extracts from custom_mask (both the eager and
-        # cuda-graph metadata paths gate the extraction on topk > 1).
-        return self.topk > 1
 
     @staticmethod
     def _host_max_seq_len(

python/sglang/srt/layers/attention/flashmla_backend.py

@@ -18,6 +18,7 @@ from sglang.kernels.ops.attention.utils import (
 )
 from sglang.kernels.ops.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend
+from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.runtime_context import get_parallel
 
@@ -102,8 +103,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
         self.cuda_graph_num_splits_view = None
         # Static K-lens buffer bound by the draft-extend graph kernel.
         self.cuda_graph_draft_extend_seq_lens_k = None
-        # Preallocated tree-mask scratch (see get_verify_buffers_to_fill_after_draft).
+        self._verify_mask = None
-        self.cuda_graph_custom_mask = None
         self._eager_kv_indices_buf = None
         # The worker fetches the tree-mask scratch from the target backend
         # only; draft-side instances must not allocate it.
@@ -285,17 +285,22 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
             self.cuda_graph_draft_extend_seq_lens_k = torch.ones(
                 max_bs, dtype=torch.int32, device="cuda"
             )
-            if not self.skip_prefill and not self.is_draft_runner:
+            # Target verify never reaches the parent's mask read: every
-                # Worst-case FULL_MASK tree-mask scratch (bool); build_tree
+            # init_forward_metadata* branch handles is_target_verify() itself and
-                # writes it in-place so the GPU-only path needs no seq_lens_sum.
+            # only falls through to super() otherwise.
-                self.cuda_graph_custom_mask = torch.zeros(
+            self._verify_mask = maybe_create_verify_mask(
-                    max_num_tokens * (self.max_context_len + self.num_draft_tokens),
+                is_draft_runner=self.is_draft_runner,
-                    dtype=torch.bool,
+                skip_prefill=self.skip_prefill,
-                    device="cuda",
+                max_bs=max_bs,
+                max_context_len=self.max_context_len,
+                num_draft_tokens=self.num_draft_tokens,
+                device=self.device,
+                is_read=False,
             )
 
-    def get_verify_buffers_to_fill_after_draft(self):
+    @property
-        return [self.cuda_graph_custom_mask, None]
+    def verify_mask(self) -> Optional[VerifyMask]:
+        return self._verify_mask
 
     def _apply_decode_target_verify_metadata(
         self,

python/sglang/srt/layers/attention/hybrid_attn_backend.py

@@ -1,4 +1,6 @@
-from typing import Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
@@ -8,6 +10,10 @@ from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
 
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.verify_mask import VerifyMask
+    from sglang.srt.speculative.spec_info import SpecInput
+
 
 class HybridAttnBackend(AttentionBackend):
     """Support different backends for prefill and decode."""
@@ -106,10 +112,17 @@ class HybridAttnBackend(AttentionBackend):
     def get_cuda_graph_seq_len_fill_value(self):
         return self.decode_backend.get_cuda_graph_seq_len_fill_value()
 
-    def target_verify_reads_custom_mask(self) -> bool:
+    @property
-        return self._select_backend(
+    def verify_mask(self) -> Optional[VerifyMask]:
+        return self._select_backend(ForwardMode.TARGET_VERIFY).verify_mask
+
+    def update_verify_buffers_to_fill_after_draft(
+        self, spec_info: SpecInput, cuda_graph_bs: Optional[int]
+    ):
+        # Plan-stream fixup goes to the same child that handed out the mask.
+        self._select_backend(
             ForwardMode.TARGET_VERIFY
-        ).target_verify_reads_custom_mask()
+        ).update_verify_buffers_to_fill_after_draft(spec_info, cuda_graph_bs)
 
     def forward(
         self,

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import logging
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import torch
 
@@ -26,6 +28,9 @@ from sglang.srt.runtime_context import get_server_args
 from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
 from sglang.srt.speculative.spec_info import SpecInput
 
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.verify_mask import VerifyMask
+
 logger = logging.getLogger(__name__)
 
 
@@ -920,16 +925,10 @@ class HybridLinearAttnBackend(AttentionBackend):
         for attn_backend in self.attn_backend_list:
             attn_backend.on_after_cuda_graph_warmup()
 
-    def get_verify_buffers_to_fill_after_draft(self):
+    @property
-        # Verify tree-mask / position buffers live on the full-attn child (the
+    def verify_mask(self) -> Optional[VerifyMask]:
-        # linear side consumes no mask). Handing them out lets the draft stage
+        # The mask lives on the full-attn child; the linear side reads none.
-        # write straight into the captured verify buffers instead of allocating
+        return self.full_attn_backend.verify_mask
-        # a fresh mask every step.
-        return self.full_attn_backend.get_verify_buffers_to_fill_after_draft()
-
-    def target_verify_reads_custom_mask(self) -> bool:
-        # Same child that hands out the mask buffer answers whether it is read.
-        return self.full_attn_backend.target_verify_reads_custom_mask()
 
     def update_verify_buffers_to_fill_after_draft(
         self, spec_info: SpecInput, cuda_graph_bs: Optional[int]

python/sglang/srt/layers/attention/tbo_backend.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Callable, List
+from typing import TYPE_CHECKING, Callable, List, Optional
 
 from sglang.srt.batch_overlap import two_batch_overlap
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 
 if TYPE_CHECKING:
+    from sglang.srt.layers.attention.verify_mask import VerifyMask
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 
@@ -41,7 +44,7 @@ class TboAttnBackend(AttentionBackend):
 
     def init_forward_metadata_out_graph(
         self,
-        forward_batch: "ForwardBatch",
+        forward_batch: ForwardBatch,
         in_capture: bool = False,
     ):
         self.primary.init_forward_metadata_out_graph(
@@ -111,7 +114,7 @@ class TboAttnBackend(AttentionBackend):
                 forward_batch=child_fb_view, in_capture=False
             )
 
-    def init_forward_metadata_in_graph(self, forward_batch: "ForwardBatch"):
+    def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
         self.primary.init_forward_metadata_in_graph(forward_batch=forward_batch)
         if not self._children_use_cuda_graph():
             return
@@ -125,7 +128,7 @@ class TboAttnBackend(AttentionBackend):
                         forward_batch=forward_batch_child
                     )
 
-    def init_forward_metadata(self, forward_batch: "ForwardBatch"):
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
         self.primary.init_forward_metadata(forward_batch=forward_batch)
         if forward_batch.tbo_children is not None:
             for child, forward_batch_child in zip(
@@ -166,9 +169,15 @@ class TboAttnBackend(AttentionBackend):
     def forward_decode(self, *args, **kwargs):
         return self.primary.forward_decode(*args, **kwargs)
 
-    def get_indexer_metadata(self, layer_id: int, forward_batch: "ForwardBatch"):
+    def get_indexer_metadata(self, layer_id: int, forward_batch: ForwardBatch):
         return self.primary.get_indexer_metadata(layer_id, forward_batch)
 
+    @property
+    def verify_mask(self) -> Optional[VerifyMask]:
+        # Needs an explicit override: the base declares this as a property, so
+        # normal lookup succeeds with None and __getattr__ below never runs.
+        return self.primary.verify_mask
+
     def __getattr__(self, name):
         # Delegate backend-specific attributes/methods not explicitly wrapped
         # above (e.g. DSV4's get_unified_swa_loc / get_swa_out_cache_loc, which

python/sglang/srt/layers/attention/triton_backend.py

@@ -21,6 +21,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
 )
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.layers.dcp import (
     cp_lse_ag_out_rs_mha,
     create_triton_kv_indices_for_dcp_triton,
@@ -296,8 +297,7 @@ class TritonAttnBackend(AttentionBackend):
             )
 
         self.forward_metadata: ForwardMetadata = None
-
+        self._verify_mask = None
-        self.cuda_graph_custom_mask = None
         # Tree-mask scratch is fetched from the target backend only.
         self.is_draft_runner = model_runner.is_draft_worker
 
@@ -496,7 +496,9 @@ class TritonAttnBackend(AttentionBackend):
                     window_kv_indices=window_kv_indices,
                 )
             )
-        custom_mask = self.cuda_graph_custom_mask
+        custom_mask = (
+            self._verify_mask.buffer if self._verify_mask is not None else None
+        )
         if (
             spec_info is not None
             and getattr(spec_info, "custom_mask", None) is not None
@@ -991,11 +993,17 @@ class TritonAttnBackend(AttentionBackend):
         else:
             self.cuda_graph_kv_indices = kv_indices_buf
 
-        if not self.skip_prefill and not self.is_draft_runner:
+        # Layout is draft * (seq_len + draft) per request (seq_mask_len cumsum
-            self.cuda_graph_custom_mask = torch.zeros(
+        # below) -- the same bound the shared sizing covers. Read as uint8.
-                (max_num_tokens * self.max_context_len),
+        self._verify_mask = maybe_create_verify_mask(
-                dtype=torch.uint8,
+            is_draft_runner=self.is_draft_runner,
+            skip_prefill=self.skip_prefill,
+            max_bs=max_bs,
+            max_context_len=self.max_context_len,
+            num_draft_tokens=self.num_draft_tokens,
             device=self.device,
+            is_read=True,
+            dtype=torch.uint8,
         )
 
         if self.sliding_window_size is not None and self.sliding_window_size > 0:
@@ -1079,8 +1087,9 @@ class TritonAttnBackend(AttentionBackend):
             )
         elif forward_mode.is_target_verify():
             custom_mask = (
-                self.cuda_graph_custom_mask
+                self._verify_mask.buffer
-                if spec_info is not None
+                if self._verify_mask is not None
+                and spec_info is not None
                 and getattr(spec_info, "custom_mask", None) is not None
                 else None
             )
@@ -1174,13 +1183,9 @@ class TritonAttnBackend(AttentionBackend):
     def get_cuda_graph_seq_len_fill_value(self):
         return 1
 
-    def get_verify_buffers_to_fill_after_draft(self):
+    @property
-        """
+    def verify_mask(self) -> Optional[VerifyMask]:
-        Return buffers for verify attention kernels that needs to be filled after draft.
+        return self._verify_mask
-
-        Typically, these are tree mask and position buffers.
-        """
-        return [self.cuda_graph_custom_mask, None]
 
     def update_verify_buffers_to_fill_after_draft(
         self, spec_info: SpecInput, cuda_graph_bs: Optional[int]

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -35,6 +35,7 @@ from sglang.srt.layers.attention.flashinfer_mla_backend import (
     FlashInferMLAAttnBackend,
     FlashInferMLAMultiStepDraftBackend,
 )
+from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
@@ -241,7 +242,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         )
 
         self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
-        self.cuda_graph_custom_mask = None
+        self._verify_mask = None
         # Tree-mask scratch is fetched from the target backend only.
         self.is_draft_runner = model_runner.is_draft_worker
 
@@ -356,19 +357,24 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 device=self.device,
             )
 
-        if self.num_draft_tokens and not self.skip_prefill and not self.is_draft_runner:
+        # Target verify never reaches the parent's mask read: it is excluded from
-            # Worst-case FULL_MASK tree-mask scratch (bool); build_tree writes it
+        # every super() dispatch (init_forward_metadata, _out_graph, forward_extend)
-            # in-place so the gpu_only path needs no seq_lens_sum.
+        # and runs the trtllm-gen kernel, which takes no mask.
-            self.cuda_graph_custom_mask = torch.zeros(
+        self._verify_mask = maybe_create_verify_mask(
-                max_num_tokens * (self.max_context_len + self.num_draft_tokens),
+            is_draft_runner=self.is_draft_runner,
-                dtype=torch.bool,
+            skip_prefill=self.skip_prefill,
+            max_bs=max_bs,
+            max_context_len=self.max_context_len,
+            num_draft_tokens=self.num_draft_tokens,
             device=self.device,
+            is_read=False,
         )
 
         super().init_cuda_graph_state(max_bs, max_num_tokens, kv_indices_buf)
 
-    def get_verify_buffers_to_fill_after_draft(self):
+    @property
-        return [self.cuda_graph_custom_mask, None]
+    def verify_mask(self) -> Optional[VerifyMask]:
+        return self._verify_mask
 
     def _init_cuda_graph_metadata(
         self,

python/sglang/srt/layers/attention/verify_mask.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import msgspec
+import torch
+
+from sglang.srt.speculative.eagle_utils import TreeMaskMode, default_tree_mask_mode
+
+
+def tree_mask_numel(
+    mode: TreeMaskMode, bs: int, num_draft_tokens: int, max_context_len: int
+) -> int:
+    """Cells the tree kernel writes for ``bs`` requests under ``mode``.
+
+    FULL_MASK reaches 100s of MB at long context; QLEN_ONLY stays in the KBs.
+    Bit-packed layouts are not sized here -- falling through to FULL_MASK would
+    over-allocate them by orders of magnitude.
+    """
+    if mode == TreeMaskMode.QLEN_ONLY:
+        per_req = num_draft_tokens * num_draft_tokens
+    elif mode == TreeMaskMode.FULL_MASK:
+        per_req = num_draft_tokens * (max_context_len + num_draft_tokens)
+    else:
+        raise NotImplementedError(f"Invalid tree mask: {mode=}")
+    return bs * per_req
+
+
+class VerifyMask(msgspec.Struct):
+    """The target-verify mask.
+
+    ``build_tree_kernel_efficient`` writes the buffer in place after draft, which
+    is what lets the worker skip the ``seq_lens_sum`` D2H sync. Keep the three
+    together: taking the buffer without its layout has the kernel write a shape
+    the reader does not expect. The kernel writes every cell even when unread, so
+    the buffer is always allocated.
+
+    Temporary home -- a phase-level buffer with no owner today (``spec_info`` is a
+    per-phase union the graph registry cannot slot).
+    """
+
+    buffer: torch.Tensor
+    mode: TreeMaskMode
+    is_read: bool = True
+
+    def fits(self, bs: int, num_draft_tokens: int) -> bool:
+        """Whether this batch's writes stay inside the buffer.
+
+        Only the compact layout is checked. FULL_MASK keeps its pre-existing
+        unconditional reuse -- its bound needs a max_context_len that composite
+        backends do not carry -- so a batch past max_bs can still overflow it
+        when draft * sum(seq_len) exceeds the buffer, as it could before.
+        """
+        if self.mode != TreeMaskMode.QLEN_ONLY:
+            return True
+        return self.buffer.numel() >= bs * num_draft_tokens * num_draft_tokens
+
+
+def maybe_create_verify_mask(
+    *,
+    is_draft_runner: bool,
+    skip_prefill: bool,
+    max_bs: int,
+    max_context_len: int,
+    num_draft_tokens: Optional[int],
+    device: torch.device | str,
+    is_read: bool,
+    dtype: torch.dtype = torch.bool,
+) -> Optional[VerifyMask]:
+    """Allocate for the captured max batch; None when nothing verifies."""
+    if is_draft_runner or skip_prefill or not num_draft_tokens:
+        return None
+    mode = default_tree_mask_mode() if is_read else TreeMaskMode.QLEN_ONLY
+    return VerifyMask(
+        buffer=torch.zeros(
+            tree_mask_numel(mode, max_bs, num_draft_tokens, max_context_len),
+            dtype=dtype,
+            device=device,
+        ),
+        mode=mode,
+        is_read=is_read,
+    )

python/sglang/srt/speculative/eagle_utils.py

@@ -153,7 +153,6 @@ def build_tree_kernel_efficient(
     num_verify_tokens: int,
     tree_mask_mode: TreeMaskMode = TreeMaskMode.FULL_MASK,
     tree_mask_buf: Optional[torch.Tensor] = None,
-    position_buf: Optional[torch.Tensor] = None,
     fill_prefix_mask: bool = True,
 ):
     draft_tokens = torch.cat((bonus_tokens.unsqueeze(1), draft_tokens), dim=1).flatten()
@@ -215,12 +214,7 @@ def build_tree_kernel_efficient(
     # position: where each token belongs to
     # e.g. if depth of each draft token is [0, 1, 1, 2] and the prompt length is 7
     # then, positions = [7, 8, 8, 9]
-    if position_buf is not None:
+    positions = torch.empty((bs * num_verify_tokens,), device=device, dtype=torch.long)
-        positions = position_buf
-    else:
-        positions = torch.empty(
-            (bs * num_verify_tokens,), device=device, dtype=torch.long
-        )
 
     if _is_npu:
         torch.ops.npu.build_tree_kernel_efficient(

python/sglang/srt/speculative/eagle_worker_common.py

@@ -342,22 +342,28 @@ def build_eagle_verify_input(
             device,
         )
 
-    # Build tree mask
+    # Write straight into the backend's buffer when it owns one and this batch
-    # Directly write to cuda graph buffers for verify attn
+    # fits; an eager batch past the captured max_bs falls back to allocating.
-    tree_mask_buf, position_buf = (
+    bs = batch.seq_lens.shape[0]
-        target_worker.model_runner.attn_backend.get_verify_buffers_to_fill_after_draft()
+    target_attn_backend = target_worker.model_runner.attn_backend
+    verify_mask = target_attn_backend.verify_mask
+    if verify_mask is None:
+        tree_mask_buf, mask_mode, fill_mask = None, tree_mask_mode, True
+    else:
+        mask_mode, fill_mask = verify_mask.mode, verify_mask.is_read
+        tree_mask_buf = (
+            verify_mask.buffer if verify_mask.fits(bs, num_draft_tokens) else None
         )
 
     # build_tree_kernel uses seq_lens_sum only to size the (non-preallocated)
-    # tree mask; over-size is safe. Skip per-iter .sum().item() D2H via UB.
+    # FULL_MASK tree mask; over-size is safe. Skip per-iter .sum().item() D2H via UB.
     seq_lens_sum = batch.seq_lens_sum
     if seq_lens_sum is None:
-        if tree_mask_buf is None:
+        if tree_mask_buf is not None or mask_mode == TreeMaskMode.QLEN_ONLY:
-            max_context_len = target_worker.model_runner.attn_backend.max_context_len
+            # Preallocated, or a QLEN_ONLY allocation sized off bs alone.
-            seq_lens_sum = batch.seq_lens.shape[0] * max_context_len
-        else:
-            # tree_mask_buf preallocated -> kernel ignores seq_lens_sum.
             seq_lens_sum = 0
+        else:
+            seq_lens_sum = bs * target_attn_backend.max_context_len
 
     (
         tree_mask,
@@ -376,10 +382,9 @@ def build_eagle_verify_input(
         topk,
         num_steps,
         num_draft_tokens,
-        tree_mask_mode,
+        mask_mode,
         tree_mask_buf,
-        position_buf,
+        fill_prefix_mask=fill_mask,
-        fill_prefix_mask=target_worker.model_runner.attn_backend.target_verify_reads_custom_mask(),
     )
 
     return EagleVerifyInput(

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -1201,9 +1201,11 @@ class EAGLEWorkerV2(BaseSpecWorker):
         retrieve_next_sibling = torch.full((bs, 1), -1, dtype=torch.long, device=device)
 
         attn_backend = self._target_worker.model_runner.attn_backend
-        mask_buf, position_buf = attn_backend.get_verify_buffers_to_fill_after_draft()
+        verify_mask = attn_backend.verify_mask
-        if mask_buf is not None:
+        # Every position in a 1-node tree is visible, so an all-True fill is
-            custom_mask = mask_buf
+        # correct under either layout.
+        if verify_mask is not None and verify_mask.fits(bs, 1):
+            custom_mask = verify_mask.buffer
             custom_mask.fill_(True)
         else:
             if batch.seq_lens_sum is not None:
@@ -1214,10 +1216,6 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 seq_lens_sum = bs * attn_backend.max_context_len
             custom_mask = torch.ones(seq_lens_sum + bs, dtype=torch.bool, device=device)
 
-        if position_buf is not None:
-            positions = position_buf
-            positions[:bs].copy_(batch.seq_lens)
-        else:
         positions = batch.seq_lens.to(torch.int64)
 
         return EagleVerifyInput(

test/registered/unit/layers/attention/test_verify_mask.py

@@ -0,0 +1,165 @@
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend
+from sglang.srt.layers.attention.verify_mask import (
+    VerifyMask,
+    maybe_create_verify_mask,
+    tree_mask_numel,
+)
+from sglang.srt.speculative.eagle_utils import TreeMaskMode, default_tree_mask_mode
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+_MAX_BS = 4
+_DRAFT = 3
+_MAX_CONTEXT_LEN = 128
+
+
+def _create(**overrides):
+    kwargs = dict(
+        is_draft_runner=False,
+        skip_prefill=False,
+        max_bs=_MAX_BS,
+        max_context_len=_MAX_CONTEXT_LEN,
+        num_draft_tokens=_DRAFT,
+        device="cpu",
+        is_read=True,
+    )
+    kwargs.update(overrides)
+    return maybe_create_verify_mask(**kwargs)
+
+
+class TestVerifyMaskSizing(CustomTestCase):
+    def test_read_mask_covers_its_layouts_write_bound(self):
+        """Whichever layout a reader gets must cover what the kernel writes:
+        FULL_MASK spans the context, QLEN_ONLY is bs * draft**2."""
+        mask = _create()
+
+        if mask.mode == TreeMaskMode.FULL_MASK:
+            bound = _MAX_BS * _DRAFT * (_MAX_CONTEXT_LEN + _DRAFT)
+        else:
+            bound = _MAX_BS * _DRAFT * _DRAFT
+        self.assertEqual(mask.mode, default_tree_mask_mode())
+        self.assertGreaterEqual(mask.buffer.numel(), bound)
+
+    def test_unread_mask_drops_the_context_dimension(self):
+        """Nothing interprets an unread layout, so it takes the compact one --
+        paying for the context dimension would be pure waste."""
+        mask = _create(is_read=False)
+
+        self.assertEqual(mask.mode, TreeMaskMode.QLEN_ONLY)
+        self.assertGreaterEqual(mask.buffer.numel(), _MAX_BS * _DRAFT * _DRAFT)
+        self.assertLess(mask.buffer.numel(), _MAX_CONTEXT_LEN)
+
+    def test_honors_dtype_override(self):
+        self.assertEqual(_create(dtype=torch.uint8).buffer.dtype, torch.uint8)
+
+
+class TestVerifyMaskCapacity(CustomTestCase):
+    """A batch past the captured max_bs must not silently reuse the compact
+    layout -- it has no context-dimension slack to absorb the overflow."""
+
+    def test_compact_layout_fits_up_to_max_bs(self):
+        # is_read=False pins QLEN_ONLY; the read layout is build-dependent.
+        mask = _create(is_read=False)
+        self.assertTrue(mask.fits(_MAX_BS, _DRAFT))
+        self.assertTrue(mask.fits(1, _DRAFT))
+
+    def test_compact_layout_does_not_fit_beyond_max_bs(self):
+        mask = _create(is_read=False)
+        self.assertFalse(mask.fits(_MAX_BS + 1, _DRAFT))
+
+    def test_full_mask_always_fits(self):
+        """FULL_MASK is exempt from the check -- see fits()."""
+        mask = VerifyMask(
+            buffer=torch.zeros(8, dtype=torch.bool), mode=TreeMaskMode.FULL_MASK
+        )
+        self.assertTrue(mask.fits(_MAX_BS * 1000, _DRAFT))
+
+
+class TestVerifyMaskGate(CustomTestCase):
+    def test_allocated_for_a_verifying_target(self):
+        self.assertIsNotNone(_create())
+
+    def test_skipped_when_nothing_verifies(self):
+        for label, overrides in (
+            ("draft runner never verifies", {"is_draft_runner": True}),
+            ("decode-only target never verifies", {"skip_prefill": True}),
+            ("no spec -> no tree", {"num_draft_tokens": None}),
+            ("zero draft tokens", {"num_draft_tokens": 0}),
+        ):
+            with self.subTest(label):
+                self.assertIsNone(_create(**overrides))
+
+
+class _FakeAttnBackend:
+    def __init__(self, verify_mask):
+        self.needs_cpu_seq_lens = False
+        self.verify_mask = verify_mask
+
+
+def _mask(numel, **kwargs):
+    return VerifyMask(
+        buffer=torch.zeros(numel, dtype=torch.bool),
+        mode=TreeMaskMode.QLEN_ONLY,
+        **kwargs,
+    )
+
+
+def _make_hybrid_backend(speculative_attention_mode, prefill_mask, decode_mask):
+    model_runner = SimpleNamespace(
+        kv_cache_dtype=None,
+        token_to_kv_pool=object(),
+        req_to_token_pool=object(),
+        server_args=SimpleNamespace(
+            speculative_attention_mode=speculative_attention_mode
+        ),
+    )
+    return HybridAttnBackend(
+        model_runner,
+        prefill_backend=_FakeAttnBackend(prefill_mask),
+        decode_backend=_FakeAttnBackend(decode_mask),
+    )
+
+
+class TestHybridAttnBackendHandsOutSelectedChildMask(CustomTestCase):
+    """Forwarding the wrong child silently falls back to a fresh mask per step."""
+
+    def test_decode_mode_uses_decode_child(self):
+        prefill_mask, decode_mask = _mask(4), _mask(8, is_read=False)
+
+        backend = _make_hybrid_backend("decode", prefill_mask, decode_mask)
+
+        self.assertIs(backend.verify_mask, decode_mask)
+
+    def test_prefill_mode_uses_prefill_child(self):
+        prefill_mask, decode_mask = _mask(4, is_read=False), _mask(8)
+
+        backend = _make_hybrid_backend("prefill", prefill_mask, decode_mask)
+
+        self.assertIs(backend.verify_mask, prefill_mask)
+
+    def test_capacity_check_needs_nothing_from_the_backend(self):
+        """A composite backend carries no max_context_len of its own: fits()
+        reaching back through the backend would raise AttributeError here."""
+        backend = _make_hybrid_backend("prefill", _mask(64, is_read=False), None)
+
+        self.assertTrue(backend.verify_mask.fits(_MAX_BS, _DRAFT))
+
+
+class TestTreeMaskNumel(CustomTestCase):
+    def test_rejects_layouts_it_cannot_size(self):
+        """A packed layout must raise, not silently take FULL_MASK's size."""
+        with self.assertRaises(NotImplementedError):
+            tree_mask_numel(
+                TreeMaskMode.QLEN_ONLY_BITPACKING, 1, _DRAFT, _MAX_CONTEXT_LEN
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/npu/attention/test_npu_ascend_backend.py

@@ -688,11 +688,9 @@ class TestGetCudaGraphSeqLenFillValue(unittest.TestCase):
 
 
 class TestGetVerifyBuffers(unittest.TestCase):
-    def test_returns_none_none(self):
+    def test_no_verify_mask(self):
         backend = object.__new__(AscendAttnBackend)
-        result = backend.get_verify_buffers_to_fill_after_draft()
+        self.assertIsNone(backend.verify_mask)
-        self.assertEqual(result, [None, None])
-        self.assertEqual(len(result), 2)
 
     def test_update_is_noop(self):
         backend = object.__new__(AscendAttnBackend)

test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py

@@ -186,7 +186,7 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase):
                 forward_batch = SimpleNamespace(forward_mode=ForwardMode.DECODE)
                 worker.draft_forward = MagicMock(return_value=graph_result)
                 attn_backend = SimpleNamespace(
-                    get_verify_buffers_to_fill_after_draft=lambda: (None, None),
+                    verify_mask=None,
                     max_context_len=1,
                 )
                 worker.target_worker = SimpleNamespace(