diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index c9fea93d6..0f6658de8 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/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] ): diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index fadc8c1a4..acc741129 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 6f78617f3..4ad2ab011 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index 2dfbd1a29..4736798a0 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index ea4947008..38de80721 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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] ): diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index ff09df9d2..6e25dd8b1 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/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=}") diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 56e645d4e..880545b50 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/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( diff --git a/test/registered/spec/utils/test_build_eagle_tree.py b/test/registered/spec/utils/test_build_eagle_tree.py index fb8d47690..b3b38a576 100644 --- a/test/registered/spec/utils/test_build_eagle_tree.py +++ b/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()