[Spec] Allocate the verify tree-mask scratch on the target backend only (#31527)

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

@@ -564,6 +564,8 @@ 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
 
     def _move_to_device(self, x: List[int]) -> torch.Tensor:
         pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True)
@@ -1496,6 +1498,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,
+            )
+
+    def get_verify_buffers_to_fill_after_draft(self):
+        return [self.cuda_graph_custom_mask, None]
 
     def replay_cuda_graph_metadata_from(
         self,

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

@@ -177,6 +177,9 @@ class FlashAttentionBackend(AttentionBackend):
         # 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
+        # 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
 
         self.use_sliding_window_kv_pool = (
             isinstance(model_runner.token_to_kv_pool, SWAKVPool)
@@ -1977,7 +1980,7 @@ class FlashAttentionBackend(AttentionBackend):
             # 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:
+            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),

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

@@ -282,6 +282,8 @@ class TritonAttnBackend(AttentionBackend):
         self.forward_metadata: ForwardMetadata = 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
 
     def get_num_kv_splits(
         self,
@@ -973,7 +975,7 @@ class TritonAttnBackend(AttentionBackend):
         else:
             self.cuda_graph_kv_indices = kv_indices_buf
 
-        if not self.skip_prefill:
+        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,

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

@@ -203,6 +203,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
 
         self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
         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
 
     def _calc_padded_blocks(self, max_seq_len: int) -> int:
         """
@@ -315,7 +317,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 device=self.device,
             )
 
-        if self.num_draft_tokens and not self.skip_prefill:
+        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(