[Perf] Skip the target-verify tree mask fill when the backend never reads it (#32886)

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

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

@@ -174,6 +174,14 @@ class AttentionBackend(ABC):
         """
         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

@@ -1524,6 +1524,10 @@ class DeepseekV4AttnBackend(
     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
+
     def replay_cuda_graph_metadata_from(
         self,
         bs: int,

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

@@ -2563,6 +2563,11 @@ class FlashAttentionBackend(AttentionBackend):
         # 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(
         seq_lens_cpu: Optional[torch.Tensor], seq_lens: torch.Tensor

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

@@ -106,6 +106,11 @@ 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(
+            ForwardMode.TARGET_VERIFY
+        ).target_verify_reads_custom_mask()
+
     def forward(
         self,
         q: Optional[torch.Tensor] = None,  # For full attention

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

@@ -927,6 +927,10 @@ class HybridLinearAttnBackend(AttentionBackend):
         # 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/speculative/eagle_utils.py

@@ -152,6 +152,7 @@ def build_tree_kernel_efficient(
     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()
 
@@ -168,7 +169,11 @@ def build_tree_kernel_efficient(
         elif tree_mask_mode == TreeMaskMode.QLEN_ONLY_BITPACKING:
             tree_mask.fill_(0)
         elif tree_mask_mode == TreeMaskMode.FULL_MASK:
-            tree_mask.fill_(True)
+            # Only the [0, seq_len) prefix columns depend on this fill; the
+            # kernel below writes every tree cell itself. Skip the (up to
+            # 100s of MB) per-step memset when nothing reads the mask.
+            if fill_prefix_mask:
+                tree_mask.fill_(True)
         else:
             raise NotImplementedError(f"Invalid tree mask: {tree_mask_mode=}")
     elif tree_mask_mode == TreeMaskMode.QLEN_ONLY:
@@ -187,13 +192,15 @@ def build_tree_kernel_efficient(
             device=device,
         )
     elif tree_mask_mode == TreeMaskMode.FULL_MASK:
-        tree_mask = torch.full(
-            (
-                seq_lens_sum * num_verify_tokens
-                + num_verify_tokens * num_verify_tokens * bs,
-            ),
-            True,
-            device=device,
+        mask_shape = (
+            seq_lens_sum * num_verify_tokens
+            + num_verify_tokens * num_verify_tokens * bs,
+        )
+        # Same reasoning as the preallocated branch above.
+        tree_mask = (
+            torch.full(mask_shape, True, dtype=torch.bool, device=device)
+            if fill_prefix_mask
+            else torch.empty(mask_shape, dtype=torch.bool, device=device)
         )
     else:
         raise NotImplementedError(f"Invalid tree mask: {tree_mask_mode=}")

python/sglang/srt/speculative/eagle_worker_common.py

@@ -379,6 +379,7 @@ def build_eagle_verify_input(
         tree_mask_mode,
         tree_mask_buf,
         position_buf,
+        fill_prefix_mask=target_worker.model_runner.attn_backend.target_verify_reads_custom_mask(),
     )
 
     return EagleVerifyInput(