From 241173724484767d151320485ca74754e36738a6 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Thu, 18 Jun 2026 13:31:12 -0700 Subject: [PATCH] [spec decoding] fully overlap spec decoding for hybrid linear attention backend (#28579) --- .../ascend_hybrid_linear_attn_backend.py | 1 + .../attention/hybrid_linear_attn_backend.py | 27 ++++++++++++++----- .../layers/attention/linear/gdn_backend.py | 2 ++ .../runner/decode_cuda_graph_runner.py | 1 + 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py index 01ec313ca..eca7ce9b9 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py @@ -130,6 +130,7 @@ class AscendMambaAttnBackendBase(MambaAttnBackendBase): forward_mode: ForwardMode, spec_info: Optional[SpecInput], seq_lens_cpu: Optional[torch.Tensor], + num_padding: Optional[int] = None, ): # out_graph passes seq_lens_cpu=None at capture; mirror the base guard. if seq_lens_cpu is None: 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 715115e58..bae807611 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -188,6 +188,9 @@ class MambaAttnBackendBase(AttentionBackend): forward_batch.forward_mode, forward_batch.spec_info, forward_batch.seq_lens_cpu if not in_capture else None, + num_padding=( + 0 if in_capture else getattr(forward_batch, "num_padding", None) + ), ) def init_forward_metadata(self, forward_batch: ForwardBatch): @@ -441,13 +444,15 @@ class MambaAttnBackendBase(AttentionBackend): forward_mode: ForwardMode, spec_info: Optional[SpecInput], seq_lens_cpu: Optional[torch.Tensor], + num_padding: Optional[int] = None, ): - if seq_lens_cpu is None: - num_padding = 0 - else: - num_padding = torch.count_nonzero( - seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value() - ) + if num_padding is None: + if seq_lens_cpu is None: + num_padding = 0 + else: + num_padding = torch.count_nonzero( + seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value() + ) # Make sure forward metadata is correctly handled for padding reqs req_pool_indices[bs - num_padding :] = 0 mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) @@ -577,6 +582,8 @@ class MambaAttnBackendBase(AttentionBackend): class Mamba2AttnBackend(MambaAttnBackendBase): """Attention backend wrapper for Mamba2Mixer kernels.""" + needs_cpu_seq_lens: bool = False + def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) config = model_runner.mamba2_config @@ -605,6 +612,9 @@ class Mamba2AttnBackend(MambaAttnBackendBase): forward_batch.forward_mode, forward_batch.spec_info, forward_batch.seq_lens_cpu if not in_capture else None, + num_padding=( + 0 if in_capture else getattr(forward_batch, "num_padding", None) + ), ) spec_info = forward_batch.spec_info draft_token_num = spec_info.draft_token_num if spec_info is not None else 1 @@ -699,6 +709,11 @@ class HybridLinearAttnBackend(AttentionBackend): # Dispatcher aliases the full-attn backend's pool refs. self.token_to_kv_pool = full_attn_backend.token_to_kv_pool self.req_to_token_pool = full_attn_backend.req_to_token_pool + self.max_context_len = getattr(full_attn_backend, "max_context_len", None) + self.needs_cpu_seq_lens = ( + full_attn_backend.needs_cpu_seq_lens + or linear_attn_backend.needs_cpu_seq_lens + ) def _is_full_attn( self, layer: Optional[RadixAttention], layer_id: Optional[int] = None diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 0194674f3..d97c09880 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -269,6 +269,8 @@ class GDNKernelDispatcher: class GDNAttnBackend(MambaAttnBackendBase): """Attention backend for GDN (Gated Delta Network) linear attention.""" + needs_cpu_seq_lens: bool = False + def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) self.conv_states_shape = ( diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index e0cc52588..72c59193f 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -154,6 +154,7 @@ def build_replay_fb_view( else forward_batch.seq_lens_sum + (bs - raw_bs) * seq_len_fill_value ), seq_lens_cpu=buffers.seq_lens_cpu[:bs], + num_padding=bs - raw_bs, encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None, out_cache_loc=getattr(forward_batch, "out_cache_loc", None), out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),