[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):
-        """
-        Return buffers of verify attention kernels that needs to be filled after draft.
-
-        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
+    @property
+    def verify_mask(self) -> Optional[VerifyMask]:
+        """The mask the draft stage fills in place, if this backend has one."""
+        return None
 
     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:
-            # DSV4's verify metadata ignores custom_mask, but handing
-            # build_tree a preallocated scratch keeps it from dynamically
-            # allocating a FULL_MASK buffer (bs * max_context_len under the
-            # GPU-only spec path) every verify step.
-            self.cuda_graph_custom_mask = torch.zeros(
-                max_num_tokens
-                * (self.max_context_len + self.speculative_num_draft_tokens),
-                dtype=torch.bool,
-                device=self.device,
-            )
+        # Verify metadata never extracts the mask. No skip_prefill notion here.
+        self._verify_mask = maybe_create_verify_mask(
+            is_draft_runner=self.is_draft_runner,
+            skip_prefill=False,
+            max_bs=max_bs,
+            max_context_len=self.max_context_len,
+            num_draft_tokens=self.speculative_num_draft_tokens,
+            device=self.device,
+            is_read=False,
+        )
 
-    def get_verify_buffers_to_fill_after_draft(self):
-        return [self.cuda_graph_custom_mask, None]
-
-    def target_verify_reads_custom_mask(self) -> bool:
-        # DSV4 verify metadata never extracts from custom_mask.
-        return False
+    @property
+    def verify_mask(self) -> Optional[VerifyMask]:
+        return self._verify_mask
 
     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),
-        # so no seq_lens_cpu / seq_lens_sum D2H sync is ever needed.
+        # tree mask is preallocated (see VerifyMask), so no
+        # 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
-        # avoid the seq_lens_sum D2H sync (see get_verify_buffers_to_fill_after_draft).
-        self.cuda_graph_custom_mask = None
+        self._verify_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,17 +2198,16 @@ class FlashAttentionBackend(AttentionBackend):
                 ),
             }
 
-            # Worst-case FULL_MASK tree-mask scratch (bool). build_tree_kernel
-            # fills it in-place, so the GPU-only path needs no seq_lens_sum.
-            # Costs max_num_tokens * max_context_len bytes (can reach 100s of
-            # MB at long context) and is fully memset every verify step.
-            if not self.skip_prefill and not self.is_draft_runner:
-                self.cuda_graph_custom_mask = torch.zeros(
-                    max_num_tokens
-                    * (self.max_context_len + self.speculative_num_draft_tokens),
-                    dtype=torch.bool,
-                    device=self.device,
-                )
+            # topk<=1 never extracts the mask; both metadata paths gate on topk > 1.
+            self._verify_mask = maybe_create_verify_mask(
+                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.speculative_num_draft_tokens,
+                device=self.device,
+                is_read=self.topk > 1,
+            )
 
             self.draft_extend_metadata = {
                 "cache_seqlens": torch.zeros(
@@ -2557,16 +2555,9 @@ class FlashAttentionBackend(AttentionBackend):
 
         return metadata, metadata_expand
 
-    def get_verify_buffers_to_fill_after_draft(self):
-        # Return the preallocated FULL_MASK tree-mask scratch so that
-        # build_tree_kernel_efficient fills it in-place and the worker never
-        # 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
+    @property
+    def verify_mask(self) -> Optional[VerifyMask]:
+        return self._verify_mask
 
     @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.cuda_graph_custom_mask = None
+        self._verify_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:
-                # Worst-case FULL_MASK tree-mask scratch (bool); build_tree
-                # writes it in-place so the GPU-only path needs no seq_lens_sum.
-                self.cuda_graph_custom_mask = torch.zeros(
-                    max_num_tokens * (self.max_context_len + self.num_draft_tokens),
-                    dtype=torch.bool,
-                    device="cuda",
-                )
+            # Target verify never reaches the parent's mask read: every
+            # init_forward_metadata* branch handles is_target_verify() itself and
+            # only falls through to super() otherwise.
+            self._verify_mask = maybe_create_verify_mask(
+                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=False,
+            )
 
-    def get_verify_buffers_to_fill_after_draft(self):
-        return [self.cuda_graph_custom_mask, None]
+    @property
+    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:
-        return self._select_backend(
+    @property
+    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):
-        # Verify tree-mask / position buffers live on the full-attn child (the
-        # linear side consumes no mask). Handing them out lets the draft stage
-        # write straight into the captured verify buffers instead of allocating
-        # 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()
+    @property
+    def verify_mask(self) -> Optional[VerifyMask]:
+        # The mask lives on the full-attn child; the linear side reads none.
+        return self.full_attn_backend.verify_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.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
 
@@ -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,12 +993,18 @@ class TritonAttnBackend(AttentionBackend):
         else:
             self.cuda_graph_kv_indices = kv_indices_buf
 
-        if not self.skip_prefill and not self.is_draft_runner:
-            self.cuda_graph_custom_mask = torch.zeros(
-                (max_num_tokens * self.max_context_len),
-                dtype=torch.uint8,
-                device=self.device,
-            )
+        # Layout is draft * (seq_len + draft) per request (seq_mask_len cumsum
+        # below) -- the same bound the shared sizing covers. Read as uint8.
+        self._verify_mask = maybe_create_verify_mask(
+            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:
             if kv_indices_buf is None:
@@ -1079,8 +1087,9 @@ class TritonAttnBackend(AttentionBackend):
             )
         elif forward_mode.is_target_verify():
             custom_mask = (
-                self.cuda_graph_custom_mask
-                if spec_info is not None
+                self._verify_mask.buffer
+                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):
-        """
-        Return buffers for verify attention kernels that needs to be filled after draft.
-
-        Typically, these are tree mask and position buffers.
-        """
-        return [self.cuda_graph_custom_mask, None]
+    @property
+    def verify_mask(self) -> Optional[VerifyMask]:
+        return self._verify_mask
 
     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:
-            # Worst-case FULL_MASK tree-mask scratch (bool); build_tree writes it
-            # in-place so the gpu_only path needs no seq_lens_sum.
-            self.cuda_graph_custom_mask = torch.zeros(
-                max_num_tokens * (self.max_context_len + self.num_draft_tokens),
-                dtype=torch.bool,
-                device=self.device,
-            )
+        # Target verify never reaches the parent's mask read: it is excluded from
+        # every super() dispatch (init_forward_metadata, _out_graph, forward_extend)
+        # and runs the trtllm-gen kernel, which takes no mask.
+        self._verify_mask = maybe_create_verify_mask(
+            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=False,
+        )
 
         super().init_cuda_graph_state(max_bs, max_num_tokens, kv_indices_buf)
 
-    def get_verify_buffers_to_fill_after_draft(self):
-        return [self.cuda_graph_custom_mask, None]
+    @property
+    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 = position_buf
-    else:
-        positions = torch.empty(
-            (bs * num_verify_tokens,), device=device, dtype=torch.long
-        )
+    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
-    # Directly write to cuda graph buffers for verify attn
-    tree_mask_buf, position_buf = (
-        target_worker.model_runner.attn_backend.get_verify_buffers_to_fill_after_draft()
-    )
+    # Write straight into the backend's buffer when it owns one and this batch
+    # fits; an eager batch past the captured max_bs falls back to allocating.
+    bs = batch.seq_lens.shape[0]
+    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:
-            max_context_len = target_worker.model_runner.attn_backend.max_context_len
-            seq_lens_sum = batch.seq_lens.shape[0] * max_context_len
-        else:
-            # tree_mask_buf preallocated -> kernel ignores seq_lens_sum.
+        if tree_mask_buf is not None or mask_mode == TreeMaskMode.QLEN_ONLY:
+            # Preallocated, or a QLEN_ONLY allocation sized off bs alone.
             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=target_worker.model_runner.attn_backend.target_verify_reads_custom_mask(),
+        fill_prefix_mask=fill_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()
-        if mask_buf is not None:
-            custom_mask = mask_buf
+        verify_mask = attn_backend.verify_mask
+        # Every position in a 1-node tree is visible, so an all-True fill is
+        # 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,11 +1216,7 @@ 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)
+        positions = batch.seq_lens.to(torch.int64)
 
         return EagleVerifyInput(
             draft_token=draft_input.bonus_tokens,