[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(

test/registered/spec/utils/test_build_eagle_tree.py

@@ -9,8 +9,8 @@ from sglang.srt.speculative.eagle_utils import (
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
-register_cuda_ci(est_time=6, stage="base-b", runner_config="1-gpu-small")
-register_amd_ci(est_time=3, suite="stage-b-test-1-gpu-small-amd")
+register_cuda_ci(est_time=8, stage="base-b", runner_config="1-gpu-small")
+register_amd_ci(est_time=4, suite="stage-b-test-1-gpu-small-amd")
 
 
 class TestBuildEagleTree(unittest.TestCase):
@@ -314,6 +314,92 @@ class TestBuildEagleTree(unittest.TestCase):
             "Draft tokens tensor does not match expected values",
         )
 
+    def test_skip_prefix_fill_preserves_tree_blocks(self):
+        """fill_prefix_mask=False must leave every kernel-written cell intact.
+
+        The fill only supplies the [0, seq_len) prefix columns; the qlen x qlen
+        tree block comes from the kernel and must be identical either way.
+        """
+        device = get_device()
+        bs, topk, spec_steps, num_draft_token = 2, 1, 3, 4
+        seq_lens = torch.tensor([5, 10], dtype=torch.int64, device=device)
+        seq_lens_sum = int(seq_lens.sum().item())
+        # topk=1 chain: token i descends from i-1; index 0 is the root.
+        parent_list = torch.tensor([[0, 0, 1]] * bs, dtype=torch.int64, device=device)
+        top_scores_index = torch.tensor(
+            [[0, 1, 2]] * bs, dtype=torch.int64, device=device
+        )
+        draft_tokens = torch.arange(
+            bs * (num_draft_token - 1), dtype=torch.int64, device=device
+        ).view(bs, -1)
+        bonus_tokens = torch.tensor([101, 102], dtype=torch.int32, device=device)
+        mask_numel = seq_lens_sum * num_draft_token + num_draft_token**2 * bs
+
+        def build(fill_prefix_mask):
+            # All-False start matches the real preallocated scratch: a skipped
+            # fill leaves the prefix stale-False.
+            tree_mask_buf = torch.zeros((mask_numel,), dtype=torch.bool, device=device)
+            return build_tree_kernel_efficient(
+                bonus_tokens=bonus_tokens,
+                parent_list=parent_list,
+                top_scores_index=top_scores_index,
+                draft_tokens=draft_tokens,
+                seq_lens=seq_lens,
+                seq_lens_sum=seq_lens_sum,
+                topk=topk,
+                spec_steps=spec_steps,
+                num_verify_tokens=num_draft_token,
+                tree_mask_buf=tree_mask_buf,
+                fill_prefix_mask=fill_prefix_mask,
+            )
+
+        def split_rows(tree_mask):
+            """Flat mask -> (all prefix columns, all tree-block cells)."""
+            prefixes, blocks = [], []
+            offset = 0
+            for seq_len in seq_lens.tolist():
+                row_len = seq_len + num_draft_token
+                for tid in range(num_draft_token):
+                    row = tree_mask[
+                        offset + row_len * tid : offset + row_len * (tid + 1)
+                    ]
+                    prefixes.append(row[:seq_len])
+                    blocks.append(row[seq_len:])
+                offset += row_len * num_draft_token
+            return torch.cat(prefixes), torch.cat(blocks)
+
+        filled = build(fill_prefix_mask=True)
+        skipped = build(fill_prefix_mask=False)
+
+        filled_prefix, filled_blocks = split_rows(filled[0])
+        skipped_prefix, skipped_blocks = split_rows(skipped[0])
+
+        self.assertTrue(
+            torch.equal(filled_blocks, skipped_blocks),
+            "Tree blocks diverged: the kernel must write every tree cell "
+            "regardless of the prefix fill",
+        )
+        # Anti-vacuous: proves the two runs really differ on the prefix.
+        self.assertTrue(filled_prefix.all(), "Fill did not mark the prefix columns")
+        self.assertFalse(
+            skipped_prefix.any(), "Skipped fill unexpectedly touched the prefix columns"
+        )
+
+        for idx, name in enumerate(
+            (
+                "positions",
+                "retrieve_index",
+                "retrieve_next_token",
+                "retrieve_next_sibling",
+                "draft_tokens",
+            ),
+            start=1,
+        ):
+            self.assertTrue(
+                torch.equal(filled[idx], skipped[idx]),
+                f"{name} diverged between filled and skipped runs",
+            )
+
 
 if __name__ == "__main__":
     unittest.main()