diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 6e3bd2efd..d70872e67 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -397,14 +397,6 @@ class AscendAttnBackend(AttentionBackend): v = layer.v_head_dim return (d == v and d in (128, 192, 256)) or (d == 192 and v == 128) - def get_verify_buffers_to_fill_after_draft(self): - """ - Return buffers for verify attention kernels that needs to be filled after draft. - - Typically, these are tree mask and position buffers. - """ - return [None, None] - 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/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 0f6658de8..fb8ced91a 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -10,6 +10,7 @@ from sglang.srt.utils.common import is_npu if TYPE_CHECKING: from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata + from sglang.srt.layers.attention.verify_mask import VerifyMask from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.speculative.spec_info import SpecInput @@ -166,21 +167,10 @@ class AttentionBackend(ABC): """ pass - def get_verify_buffers_to_fill_after_draft(self): - """ - Return buffers of verify attention kernels that needs to be filled after draft. - - Typically, these are tree mask and position buffers. - """ - 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 + @property + def verify_mask(self) -> Optional[VerifyMask]: + """The mask the draft stage fills in place, if this backend has one.""" + return None 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 acc741129..bd4385dfa 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -57,6 +57,7 @@ from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( SparsePrefillChunkCache, SparsePrefillWorkspace, ) +from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_parallel @@ -571,7 +572,7 @@ 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 + self._verify_mask = None def _move_to_device(self, x: List[int]) -> torch.Tensor: pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) @@ -1509,24 +1510,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, - ) + # Verify metadata never extracts the mask. No skip_prefill notion here. + self._verify_mask = maybe_create_verify_mask( + is_draft_runner=self.is_draft_runner, + skip_prefill=False, + max_bs=max_bs, + max_context_len=self.max_context_len, + num_draft_tokens=self.speculative_num_draft_tokens, + device=self.device, + is_read=False, + ) - 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 + @property + def verify_mask(self) -> Optional[VerifyMask]: + return self._verify_mask 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 4ad2ab011..7448e8d3c 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -18,6 +18,7 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import ( ) from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy from sglang.srt.layers.cp.utils import is_cp_v2_active from sglang.srt.layers.radix_attention import AttentionType @@ -179,15 +180,13 @@ class FlashAttentionBackend(AttentionBackend): self.max_context_len + self.page_size - 1 ) // self.page_size # Page table is built on-device (build_trtllm_mha_page_table) and the - # tree-mask scratch is preallocated (get_verify_buffers_to_fill_after_draft), - # so no seq_lens_cpu / seq_lens_sum D2H sync is ever needed. + # tree mask is preallocated (see VerifyMask), so no + # seq_lens_cpu / seq_lens_sum D2H sync is ever needed. self.needs_cpu_seq_lens = False self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA self.skip_prefill = skip_prefill self.attn_cp_size = model_runner.ps.attn_cp_size - # 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 + self._verify_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 @@ -2199,17 +2198,16 @@ class FlashAttentionBackend(AttentionBackend): ), } - # Worst-case FULL_MASK tree-mask scratch (bool). build_tree_kernel - # 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 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), - dtype=torch.bool, - device=self.device, - ) + # topk<=1 never extracts the mask; both metadata paths gate on topk > 1. + self._verify_mask = maybe_create_verify_mask( + is_draft_runner=self.is_draft_runner, + skip_prefill=self.skip_prefill, + max_bs=max_bs, + max_context_len=self.max_context_len, + num_draft_tokens=self.speculative_num_draft_tokens, + device=self.device, + is_read=self.topk > 1, + ) self.draft_extend_metadata = { "cache_seqlens": torch.zeros( @@ -2557,16 +2555,9 @@ class FlashAttentionBackend(AttentionBackend): return metadata, metadata_expand - def get_verify_buffers_to_fill_after_draft(self): - # Return the preallocated FULL_MASK tree-mask scratch so that - # build_tree_kernel_efficient fills it in-place and the worker never - # 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 + @property + def verify_mask(self) -> Optional[VerifyMask]: + return self._verify_mask @staticmethod def _host_max_seq_len( diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index 1abc64f84..ce628b818 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -18,6 +18,7 @@ from sglang.kernels.ops.attention.utils import ( ) from sglang.kernels.ops.quantization.fp8_kernel import scaled_fp8_quant from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend +from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_parallel @@ -102,8 +103,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend): self.cuda_graph_num_splits_view = None # Static K-lens buffer bound by the draft-extend graph kernel. self.cuda_graph_draft_extend_seq_lens_k = None - # Preallocated tree-mask scratch (see get_verify_buffers_to_fill_after_draft). - self.cuda_graph_custom_mask = None + self._verify_mask = None self._eager_kv_indices_buf = None # The worker fetches the tree-mask scratch from the target backend # only; draft-side instances must not allocate it. @@ -285,17 +285,22 @@ class FlashMLABackend(FlashInferMLAAttnBackend): self.cuda_graph_draft_extend_seq_lens_k = torch.ones( max_bs, dtype=torch.int32, device="cuda" ) - if 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( - max_num_tokens * (self.max_context_len + self.num_draft_tokens), - dtype=torch.bool, - device="cuda", - ) + # Target verify never reaches the parent's mask read: every + # init_forward_metadata* branch handles is_target_verify() itself and + # only falls through to super() otherwise. + self._verify_mask = maybe_create_verify_mask( + is_draft_runner=self.is_draft_runner, + skip_prefill=self.skip_prefill, + max_bs=max_bs, + max_context_len=self.max_context_len, + num_draft_tokens=self.num_draft_tokens, + device=self.device, + is_read=False, + ) - def get_verify_buffers_to_fill_after_draft(self): - return [self.cuda_graph_custom_mask, None] + @property + def verify_mask(self) -> Optional[VerifyMask]: + return self._verify_mask def _apply_decode_target_verify_metadata( self, diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index 4736798a0..16ce62cea 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_attn_backend.py @@ -1,4 +1,6 @@ -from typing import Optional +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional import torch @@ -8,6 +10,10 @@ from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.model_runner import ModelRunner +if TYPE_CHECKING: + from sglang.srt.layers.attention.verify_mask import VerifyMask + from sglang.srt.speculative.spec_info import SpecInput + class HybridAttnBackend(AttentionBackend): """Support different backends for prefill and decode.""" @@ -106,10 +112,17 @@ 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( + @property + def verify_mask(self) -> Optional[VerifyMask]: + return self._select_backend(ForwardMode.TARGET_VERIFY).verify_mask + + def update_verify_buffers_to_fill_after_draft( + self, spec_info: SpecInput, cuda_graph_bs: Optional[int] + ): + # Plan-stream fixup goes to the same child that handed out the mask. + self._select_backend( ForwardMode.TARGET_VERIFY - ).target_verify_reads_custom_mask() + ).update_verify_buffers_to_fill_after_draft(spec_info, cuda_graph_bs) def forward( self, 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 38de80721..b335700dd 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -1,5 +1,7 @@ +from __future__ import annotations + import logging -from typing import Optional, Union +from typing import TYPE_CHECKING, Optional, Union import torch @@ -26,6 +28,9 @@ from sglang.srt.runtime_context import get_server_args from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput from sglang.srt.speculative.spec_info import SpecInput +if TYPE_CHECKING: + from sglang.srt.layers.attention.verify_mask import VerifyMask + logger = logging.getLogger(__name__) @@ -920,16 +925,10 @@ class HybridLinearAttnBackend(AttentionBackend): for attn_backend in self.attn_backend_list: attn_backend.on_after_cuda_graph_warmup() - def get_verify_buffers_to_fill_after_draft(self): - # Verify tree-mask / position buffers live on the full-attn child (the - # linear side consumes no mask). Handing them out lets the draft stage - # write straight into the captured verify buffers instead of allocating - # 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() + @property + def verify_mask(self) -> Optional[VerifyMask]: + # The mask lives on the full-attn child; the linear side reads none. + return self.full_attn_backend.verify_mask 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/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index d595f3520..27867541d 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -1,10 +1,13 @@ +from __future__ import annotations + from types import SimpleNamespace -from typing import TYPE_CHECKING, Callable, List +from typing import TYPE_CHECKING, Callable, List, Optional from sglang.srt.batch_overlap import two_batch_overlap from sglang.srt.layers.attention.base_attn_backend import AttentionBackend if TYPE_CHECKING: + from sglang.srt.layers.attention.verify_mask import VerifyMask from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -41,7 +44,7 @@ class TboAttnBackend(AttentionBackend): def init_forward_metadata_out_graph( self, - forward_batch: "ForwardBatch", + forward_batch: ForwardBatch, in_capture: bool = False, ): self.primary.init_forward_metadata_out_graph( @@ -111,7 +114,7 @@ class TboAttnBackend(AttentionBackend): forward_batch=child_fb_view, in_capture=False ) - def init_forward_metadata_in_graph(self, forward_batch: "ForwardBatch"): + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): self.primary.init_forward_metadata_in_graph(forward_batch=forward_batch) if not self._children_use_cuda_graph(): return @@ -125,7 +128,7 @@ class TboAttnBackend(AttentionBackend): forward_batch=forward_batch_child ) - def init_forward_metadata(self, forward_batch: "ForwardBatch"): + def init_forward_metadata(self, forward_batch: ForwardBatch): self.primary.init_forward_metadata(forward_batch=forward_batch) if forward_batch.tbo_children is not None: for child, forward_batch_child in zip( @@ -166,9 +169,15 @@ class TboAttnBackend(AttentionBackend): def forward_decode(self, *args, **kwargs): return self.primary.forward_decode(*args, **kwargs) - def get_indexer_metadata(self, layer_id: int, forward_batch: "ForwardBatch"): + def get_indexer_metadata(self, layer_id: int, forward_batch: ForwardBatch): return self.primary.get_indexer_metadata(layer_id, forward_batch) + @property + def verify_mask(self) -> Optional[VerifyMask]: + # Needs an explicit override: the base declares this as a property, so + # normal lookup succeeds with None and __getattr__ below never runs. + return self.primary.verify_mask + def __getattr__(self, name): # Delegate backend-specific attributes/methods not explicitly wrapped # above (e.g. DSV4's get_unified_swa_loc / get_swa_out_cache_loc, which diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index e5db20875..02fe634b0 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -21,6 +21,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( ) from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask from sglang.srt.layers.dcp import ( cp_lse_ag_out_rs_mha, create_triton_kv_indices_for_dcp_triton, @@ -296,8 +297,7 @@ class TritonAttnBackend(AttentionBackend): ) self.forward_metadata: ForwardMetadata = None - - self.cuda_graph_custom_mask = None + self._verify_mask = None # Tree-mask scratch is fetched from the target backend only. self.is_draft_runner = model_runner.is_draft_worker @@ -496,7 +496,9 @@ class TritonAttnBackend(AttentionBackend): window_kv_indices=window_kv_indices, ) ) - custom_mask = self.cuda_graph_custom_mask + custom_mask = ( + self._verify_mask.buffer if self._verify_mask is not None else None + ) if ( spec_info is not None and getattr(spec_info, "custom_mask", None) is not None @@ -991,12 +993,18 @@ class TritonAttnBackend(AttentionBackend): else: self.cuda_graph_kv_indices = kv_indices_buf - 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, - device=self.device, - ) + # Layout is draft * (seq_len + draft) per request (seq_mask_len cumsum + # below) -- the same bound the shared sizing covers. Read as uint8. + self._verify_mask = maybe_create_verify_mask( + is_draft_runner=self.is_draft_runner, + skip_prefill=self.skip_prefill, + max_bs=max_bs, + max_context_len=self.max_context_len, + num_draft_tokens=self.num_draft_tokens, + device=self.device, + is_read=True, + dtype=torch.uint8, + ) if self.sliding_window_size is not None and self.sliding_window_size > 0: if kv_indices_buf is None: @@ -1079,8 +1087,9 @@ class TritonAttnBackend(AttentionBackend): ) elif forward_mode.is_target_verify(): custom_mask = ( - self.cuda_graph_custom_mask - if spec_info is not None + self._verify_mask.buffer + if self._verify_mask is not None + and spec_info is not None and getattr(spec_info, "custom_mask", None) is not None else None ) @@ -1174,13 +1183,9 @@ class TritonAttnBackend(AttentionBackend): def get_cuda_graph_seq_len_fill_value(self): return 1 - def get_verify_buffers_to_fill_after_draft(self): - """ - Return buffers for verify attention kernels that needs to be filled after draft. - - Typically, these are tree mask and position buffers. - """ - return [self.cuda_graph_custom_mask, None] + @property + def verify_mask(self) -> Optional[VerifyMask]: + return self._verify_mask 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/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index a74889c55..1c1d07401 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -35,6 +35,7 @@ from sglang.srt.layers.attention.flashinfer_mla_backend import ( FlashInferMLAAttnBackend, FlashInferMLAMultiStepDraftBackend, ) +from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, @@ -241,7 +242,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ) self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens - self.cuda_graph_custom_mask = None + self._verify_mask = None # Tree-mask scratch is fetched from the target backend only. self.is_draft_runner = model_runner.is_draft_worker @@ -356,19 +357,24 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): device=self.device, ) - 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( - max_num_tokens * (self.max_context_len + self.num_draft_tokens), - dtype=torch.bool, - device=self.device, - ) + # Target verify never reaches the parent's mask read: it is excluded from + # every super() dispatch (init_forward_metadata, _out_graph, forward_extend) + # and runs the trtllm-gen kernel, which takes no mask. + self._verify_mask = maybe_create_verify_mask( + is_draft_runner=self.is_draft_runner, + skip_prefill=self.skip_prefill, + max_bs=max_bs, + max_context_len=self.max_context_len, + num_draft_tokens=self.num_draft_tokens, + device=self.device, + is_read=False, + ) super().init_cuda_graph_state(max_bs, max_num_tokens, kv_indices_buf) - def get_verify_buffers_to_fill_after_draft(self): - return [self.cuda_graph_custom_mask, None] + @property + def verify_mask(self) -> Optional[VerifyMask]: + return self._verify_mask def _init_cuda_graph_metadata( self, diff --git a/python/sglang/srt/layers/attention/verify_mask.py b/python/sglang/srt/layers/attention/verify_mask.py new file mode 100644 index 000000000..417c22882 --- /dev/null +++ b/python/sglang/srt/layers/attention/verify_mask.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from typing import Optional + +import msgspec +import torch + +from sglang.srt.speculative.eagle_utils import TreeMaskMode, default_tree_mask_mode + + +def tree_mask_numel( + mode: TreeMaskMode, bs: int, num_draft_tokens: int, max_context_len: int +) -> int: + """Cells the tree kernel writes for ``bs`` requests under ``mode``. + + FULL_MASK reaches 100s of MB at long context; QLEN_ONLY stays in the KBs. + Bit-packed layouts are not sized here -- falling through to FULL_MASK would + over-allocate them by orders of magnitude. + """ + if mode == TreeMaskMode.QLEN_ONLY: + per_req = num_draft_tokens * num_draft_tokens + elif mode == TreeMaskMode.FULL_MASK: + per_req = num_draft_tokens * (max_context_len + num_draft_tokens) + else: + raise NotImplementedError(f"Invalid tree mask: {mode=}") + return bs * per_req + + +class VerifyMask(msgspec.Struct): + """The target-verify mask. + + ``build_tree_kernel_efficient`` writes the buffer in place after draft, which + is what lets the worker skip the ``seq_lens_sum`` D2H sync. Keep the three + together: taking the buffer without its layout has the kernel write a shape + the reader does not expect. The kernel writes every cell even when unread, so + the buffer is always allocated. + + Temporary home -- a phase-level buffer with no owner today (``spec_info`` is a + per-phase union the graph registry cannot slot). + """ + + buffer: torch.Tensor + mode: TreeMaskMode + is_read: bool = True + + def fits(self, bs: int, num_draft_tokens: int) -> bool: + """Whether this batch's writes stay inside the buffer. + + Only the compact layout is checked. FULL_MASK keeps its pre-existing + unconditional reuse -- its bound needs a max_context_len that composite + backends do not carry -- so a batch past max_bs can still overflow it + when draft * sum(seq_len) exceeds the buffer, as it could before. + """ + if self.mode != TreeMaskMode.QLEN_ONLY: + return True + return self.buffer.numel() >= bs * num_draft_tokens * num_draft_tokens + + +def maybe_create_verify_mask( + *, + is_draft_runner: bool, + skip_prefill: bool, + max_bs: int, + max_context_len: int, + num_draft_tokens: Optional[int], + device: torch.device | str, + is_read: bool, + dtype: torch.dtype = torch.bool, +) -> Optional[VerifyMask]: + """Allocate for the captured max batch; None when nothing verifies.""" + if is_draft_runner or skip_prefill or not num_draft_tokens: + return None + mode = default_tree_mask_mode() if is_read else TreeMaskMode.QLEN_ONLY + return VerifyMask( + buffer=torch.zeros( + tree_mask_numel(mode, max_bs, num_draft_tokens, max_context_len), + dtype=dtype, + device=device, + ), + mode=mode, + is_read=is_read, + ) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index c9037aaf1..078d0d1a9 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -153,7 +153,6 @@ def build_tree_kernel_efficient( num_verify_tokens: int, 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() @@ -215,12 +214,7 @@ def build_tree_kernel_efficient( # position: where each token belongs to # e.g. if depth of each draft token is [0, 1, 1, 2] and the prompt length is 7 # then, positions = [7, 8, 8, 9] - if position_buf is not None: - positions = position_buf - else: - positions = torch.empty( - (bs * num_verify_tokens,), device=device, dtype=torch.long - ) + positions = torch.empty((bs * num_verify_tokens,), device=device, dtype=torch.long) if _is_npu: torch.ops.npu.build_tree_kernel_efficient( diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 880545b50..efb3891a6 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/python/sglang/srt/speculative/eagle_worker_common.py @@ -342,22 +342,28 @@ def build_eagle_verify_input( device, ) - # Build tree mask - # Directly write to cuda graph buffers for verify attn - tree_mask_buf, position_buf = ( - target_worker.model_runner.attn_backend.get_verify_buffers_to_fill_after_draft() - ) + # Write straight into the backend's buffer when it owns one and this batch + # fits; an eager batch past the captured max_bs falls back to allocating. + bs = batch.seq_lens.shape[0] + target_attn_backend = target_worker.model_runner.attn_backend + verify_mask = target_attn_backend.verify_mask + if verify_mask is None: + tree_mask_buf, mask_mode, fill_mask = None, tree_mask_mode, True + else: + mask_mode, fill_mask = verify_mask.mode, verify_mask.is_read + tree_mask_buf = ( + verify_mask.buffer if verify_mask.fits(bs, num_draft_tokens) else None + ) # build_tree_kernel uses seq_lens_sum only to size the (non-preallocated) - # tree mask; over-size is safe. Skip per-iter .sum().item() D2H via UB. + # FULL_MASK tree mask; over-size is safe. Skip per-iter .sum().item() D2H via UB. seq_lens_sum = batch.seq_lens_sum if seq_lens_sum is None: - if tree_mask_buf is None: - max_context_len = target_worker.model_runner.attn_backend.max_context_len - seq_lens_sum = batch.seq_lens.shape[0] * max_context_len - else: - # tree_mask_buf preallocated -> kernel ignores seq_lens_sum. + if tree_mask_buf is not None or mask_mode == TreeMaskMode.QLEN_ONLY: + # Preallocated, or a QLEN_ONLY allocation sized off bs alone. seq_lens_sum = 0 + else: + seq_lens_sum = bs * target_attn_backend.max_context_len ( tree_mask, @@ -376,10 +382,9 @@ def build_eagle_verify_input( topk, num_steps, num_draft_tokens, - tree_mask_mode, + mask_mode, tree_mask_buf, - position_buf, - fill_prefix_mask=target_worker.model_runner.attn_backend.target_verify_reads_custom_mask(), + fill_prefix_mask=fill_mask, ) return EagleVerifyInput( diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index ee4357578..047260130 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -1201,9 +1201,11 @@ class EAGLEWorkerV2(BaseSpecWorker): retrieve_next_sibling = torch.full((bs, 1), -1, dtype=torch.long, device=device) attn_backend = self._target_worker.model_runner.attn_backend - mask_buf, position_buf = attn_backend.get_verify_buffers_to_fill_after_draft() - if mask_buf is not None: - custom_mask = mask_buf + verify_mask = attn_backend.verify_mask + # Every position in a 1-node tree is visible, so an all-True fill is + # correct under either layout. + if verify_mask is not None and verify_mask.fits(bs, 1): + custom_mask = verify_mask.buffer custom_mask.fill_(True) else: if batch.seq_lens_sum is not None: @@ -1214,11 +1216,7 @@ class EAGLEWorkerV2(BaseSpecWorker): seq_lens_sum = bs * attn_backend.max_context_len custom_mask = torch.ones(seq_lens_sum + bs, dtype=torch.bool, device=device) - if position_buf is not None: - positions = position_buf - positions[:bs].copy_(batch.seq_lens) - else: - positions = batch.seq_lens.to(torch.int64) + positions = batch.seq_lens.to(torch.int64) return EagleVerifyInput( draft_token=draft_input.bonus_tokens, diff --git a/test/registered/unit/layers/attention/test_verify_mask.py b/test/registered/unit/layers/attention/test_verify_mask.py new file mode 100644 index 000000000..b16331d95 --- /dev/null +++ b/test/registered/unit/layers/attention/test_verify_mask.py @@ -0,0 +1,165 @@ +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend +from sglang.srt.layers.attention.verify_mask import ( + VerifyMask, + maybe_create_verify_mask, + tree_mask_numel, +) +from sglang.srt.speculative.eagle_utils import TreeMaskMode, default_tree_mask_mode +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +_MAX_BS = 4 +_DRAFT = 3 +_MAX_CONTEXT_LEN = 128 + + +def _create(**overrides): + kwargs = dict( + is_draft_runner=False, + skip_prefill=False, + max_bs=_MAX_BS, + max_context_len=_MAX_CONTEXT_LEN, + num_draft_tokens=_DRAFT, + device="cpu", + is_read=True, + ) + kwargs.update(overrides) + return maybe_create_verify_mask(**kwargs) + + +class TestVerifyMaskSizing(CustomTestCase): + def test_read_mask_covers_its_layouts_write_bound(self): + """Whichever layout a reader gets must cover what the kernel writes: + FULL_MASK spans the context, QLEN_ONLY is bs * draft**2.""" + mask = _create() + + if mask.mode == TreeMaskMode.FULL_MASK: + bound = _MAX_BS * _DRAFT * (_MAX_CONTEXT_LEN + _DRAFT) + else: + bound = _MAX_BS * _DRAFT * _DRAFT + self.assertEqual(mask.mode, default_tree_mask_mode()) + self.assertGreaterEqual(mask.buffer.numel(), bound) + + def test_unread_mask_drops_the_context_dimension(self): + """Nothing interprets an unread layout, so it takes the compact one -- + paying for the context dimension would be pure waste.""" + mask = _create(is_read=False) + + self.assertEqual(mask.mode, TreeMaskMode.QLEN_ONLY) + self.assertGreaterEqual(mask.buffer.numel(), _MAX_BS * _DRAFT * _DRAFT) + self.assertLess(mask.buffer.numel(), _MAX_CONTEXT_LEN) + + def test_honors_dtype_override(self): + self.assertEqual(_create(dtype=torch.uint8).buffer.dtype, torch.uint8) + + +class TestVerifyMaskCapacity(CustomTestCase): + """A batch past the captured max_bs must not silently reuse the compact + layout -- it has no context-dimension slack to absorb the overflow.""" + + def test_compact_layout_fits_up_to_max_bs(self): + # is_read=False pins QLEN_ONLY; the read layout is build-dependent. + mask = _create(is_read=False) + self.assertTrue(mask.fits(_MAX_BS, _DRAFT)) + self.assertTrue(mask.fits(1, _DRAFT)) + + def test_compact_layout_does_not_fit_beyond_max_bs(self): + mask = _create(is_read=False) + self.assertFalse(mask.fits(_MAX_BS + 1, _DRAFT)) + + def test_full_mask_always_fits(self): + """FULL_MASK is exempt from the check -- see fits().""" + mask = VerifyMask( + buffer=torch.zeros(8, dtype=torch.bool), mode=TreeMaskMode.FULL_MASK + ) + self.assertTrue(mask.fits(_MAX_BS * 1000, _DRAFT)) + + +class TestVerifyMaskGate(CustomTestCase): + def test_allocated_for_a_verifying_target(self): + self.assertIsNotNone(_create()) + + def test_skipped_when_nothing_verifies(self): + for label, overrides in ( + ("draft runner never verifies", {"is_draft_runner": True}), + ("decode-only target never verifies", {"skip_prefill": True}), + ("no spec -> no tree", {"num_draft_tokens": None}), + ("zero draft tokens", {"num_draft_tokens": 0}), + ): + with self.subTest(label): + self.assertIsNone(_create(**overrides)) + + +class _FakeAttnBackend: + def __init__(self, verify_mask): + self.needs_cpu_seq_lens = False + self.verify_mask = verify_mask + + +def _mask(numel, **kwargs): + return VerifyMask( + buffer=torch.zeros(numel, dtype=torch.bool), + mode=TreeMaskMode.QLEN_ONLY, + **kwargs, + ) + + +def _make_hybrid_backend(speculative_attention_mode, prefill_mask, decode_mask): + model_runner = SimpleNamespace( + kv_cache_dtype=None, + token_to_kv_pool=object(), + req_to_token_pool=object(), + server_args=SimpleNamespace( + speculative_attention_mode=speculative_attention_mode + ), + ) + return HybridAttnBackend( + model_runner, + prefill_backend=_FakeAttnBackend(prefill_mask), + decode_backend=_FakeAttnBackend(decode_mask), + ) + + +class TestHybridAttnBackendHandsOutSelectedChildMask(CustomTestCase): + """Forwarding the wrong child silently falls back to a fresh mask per step.""" + + def test_decode_mode_uses_decode_child(self): + prefill_mask, decode_mask = _mask(4), _mask(8, is_read=False) + + backend = _make_hybrid_backend("decode", prefill_mask, decode_mask) + + self.assertIs(backend.verify_mask, decode_mask) + + def test_prefill_mode_uses_prefill_child(self): + prefill_mask, decode_mask = _mask(4, is_read=False), _mask(8) + + backend = _make_hybrid_backend("prefill", prefill_mask, decode_mask) + + self.assertIs(backend.verify_mask, prefill_mask) + + def test_capacity_check_needs_nothing_from_the_backend(self): + """A composite backend carries no max_context_len of its own: fits() + reaching back through the backend would raise AttributeError here.""" + backend = _make_hybrid_backend("prefill", _mask(64, is_read=False), None) + + self.assertTrue(backend.verify_mask.fits(_MAX_BS, _DRAFT)) + + +class TestTreeMaskNumel(CustomTestCase): + def test_rejects_layouts_it_cannot_size(self): + """A packed layout must raise, not silently take FULL_MASK's size.""" + with self.assertRaises(NotImplementedError): + tree_mask_numel( + TreeMaskMode.QLEN_ONLY_BITPACKING, 1, _DRAFT, _MAX_CONTEXT_LEN + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/npu/attention/test_npu_ascend_backend.py b/test/registered/unit/npu/attention/test_npu_ascend_backend.py index e9955cd77..867a4a38a 100644 --- a/test/registered/unit/npu/attention/test_npu_ascend_backend.py +++ b/test/registered/unit/npu/attention/test_npu_ascend_backend.py @@ -688,11 +688,9 @@ class TestGetCudaGraphSeqLenFillValue(unittest.TestCase): class TestGetVerifyBuffers(unittest.TestCase): - def test_returns_none_none(self): + def test_no_verify_mask(self): backend = object.__new__(AscendAttnBackend) - result = backend.get_verify_buffers_to_fill_after_draft() - self.assertEqual(result, [None, None]) - self.assertEqual(len(result), 2) + self.assertIsNone(backend.verify_mask) def test_update_is_noop(self): backend = object.__new__(AscendAttnBackend) diff --git a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py index 3f42d3c3b..addccff00 100644 --- a/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py +++ b/test/registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py @@ -186,7 +186,7 @@ class TestEagleWorkerV2BackendFallback(CustomTestCase): forward_batch = SimpleNamespace(forward_mode=ForwardMode.DECODE) worker.draft_forward = MagicMock(return_value=graph_result) attn_backend = SimpleNamespace( - get_verify_buffers_to_fill_after_draft=lambda: (None, None), + verify_mask=None, max_context_len=1, ) worker.target_worker = SimpleNamespace(