diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 7dc34d158..6467c682d 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index a6c0122d5..336ffc6ef 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/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), diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 55091f114..5b60bb40a 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index aeec6d1e4..6e6d6fce4 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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(