From 3b9db3a1f054398aaa07ab5ab98c2085b8807452 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Thu, 18 Jun 2026 11:42:31 +0800 Subject: [PATCH] [Mamba][GDN] Deduplicate spec conv-window intermediate cache via sliding window layout (#28302) Co-authored-by: luoyuan.luo --- python/sglang/srt/disaggregation/decode.py | 2 + .../attention/hybrid_linear_attn_backend.py | 7 +- .../mamba/mamba_state_scatter_triton.py | 153 ++++++++++++++++++ python/sglang/srt/mem_cache/memory_pool.py | 139 +++++++++++++--- .../model_runner_kv_cache_mixin.py | 2 + 5 files changed, 283 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 6fb6aa2f6..1b792e0ff 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -195,6 +195,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): enable_overlap_schedule: bool, mamba_size: int = None, start_layer: int = None, + speculative_eagle_topk: Optional[int] = None, ): DecodeReqToTokenPool.__init__( self, @@ -238,6 +239,7 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): device=device, enable_mamba_extra_buffer=self.enable_mamba_extra_buffer, speculative_num_draft_tokens=speculative_num_draft_tokens, + speculative_eagle_topk=speculative_eagle_topk, ) def clear(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 19ecb95cc..715115e58 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -11,6 +11,7 @@ from sglang.srt.layers.attention.mamba.mamba2_metadata import ( Mamba2Metadata, ) from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import ( + fused_conv_window_scatter_with_mask, fused_mamba_state_scatter_with_mask, track_mamba_states_if_needed, ) @@ -922,7 +923,9 @@ class HybridLinearAttnBackend(AttentionBackend): state_indices_tensor, last_correct_step_indices, ) - fused_mamba_state_scatter_with_mask( + # conv intermediate uses the deduplicated sliding-window (overlapping) + # layout, so it needs the strided-read scatter variant. + fused_conv_window_scatter_with_mask( conv_states, intermediate_conv_window_cache, state_indices_tensor, @@ -939,7 +942,7 @@ class HybridLinearAttnBackend(AttentionBackend): mamba_track_indices, mamba_steps_to_track, ) - fused_mamba_state_scatter_with_mask( + fused_conv_window_scatter_with_mask( conv_states, intermediate_conv_window_cache, mamba_track_indices, diff --git a/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py b/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py index dcdc83e4a..9a8ec50d5 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py +++ b/python/sglang/srt/layers/attention/mamba/mamba_state_scatter_triton.py @@ -286,3 +286,156 @@ def fused_mamba_state_scatter_with_mask( dst_req_size, BLOCK_SIZE=BLOCK_SIZE, ) + + +@triton.jit +def _fused_conv_window_scatter_with_mask_kernel( + src_ptr, + dst_ptr, + dst_indices_raw_ptr, # [total_requests] + step_indices_raw_ptr, # [total_requests], entry >= 0 means valid + elem_per_entry: tl.constexpr, # dim * (K-1) + KM1: tl.constexpr, # K-1 (conv window width) + src_layer_stride, + src_req_stride, + src_step_stride, + src_dim_stride, + src_win_stride, + dst_layer_stride, + dst_req_stride, + src_req_size, + src_step_size, + dst_req_size, + BLOCK_SIZE: tl.constexpr, +): + """Scatter accepted conv windows from the deduplicated sliding-window source. + + Unlike ``_fused_mamba_state_scatter_with_mask_kernel`` (which flat-copies a + contiguous per-step state row), the source here is an *overlapping* view: the + deduplicated layout keeps one shared ``[dim, D+K-2]`` buffer per (layer, slot) + and step ``t``'s window is the slice ``shared[:, t:t+K-1]``. That window is + non-contiguous, so we index every ``(dim, win)`` element through the view's + strides (``src_step_stride`` / ``src_dim_stride`` / ``src_win_stride``). The + destination conv-state row stays contiguous in ``(dim, K-1)`` order. + """ + pid_req = tl.program_id(0) + pid_layer = tl.program_id(1).to(tl.int64) + pid_block = tl.program_id(2).to(tl.int64) + + step_idx = tl.load(step_indices_raw_ptr + pid_req).to(tl.int64) + if step_idx < 0: + return + + dst_idx = tl.load(dst_indices_raw_ptr + pid_req).to(tl.int64) + src_idx = pid_req + + if not ( + (dst_idx >= 0) + & (dst_idx < dst_req_size) + & (src_idx < src_req_size) + & (step_idx < src_step_size) + ): + return + + start = pid_block * BLOCK_SIZE + e = start + tl.arange(0, BLOCK_SIZE) + mask = e < elem_per_entry + + # Decode the flat (dim, K-1)-row element index into (dim, win) coordinates. + d = e // KM1 + w = e % KM1 + + src_off = ( + pid_layer * src_layer_stride + + src_idx * src_req_stride + + step_idx * src_step_stride + + d * src_dim_stride + + w * src_win_stride + ) + # dst window is contiguous in (dim, K-1) order -> flat element index `e`. + dst_off = pid_layer * dst_layer_stride + dst_idx * dst_req_stride + e + + data = tl.load(src_ptr + src_off, mask=mask, other=0.0) + tl.store(dst_ptr + dst_off, data, mask=mask) + + +def fused_conv_window_scatter_with_mask( + dst: torch.Tensor, # conv_states [num_layers, cache_size, dim, K-1] (contiguous) + src: torch.Tensor, # deduped conv-window view [num_layers, spec_size, draft_tokens, dim, K-1] + dst_indices_raw: torch.Tensor, # [total_requests] + step_indices_raw: torch.Tensor, # [total_requests], entry >= 0 means valid +): + """Conv-window variant of :func:`fused_mamba_state_scatter_with_mask`. + + ``src`` is the deduplicated sliding-window conv-intermediate cache: an + overlapping ``as_strided`` view over a shared ``[..., dim, D+K-2]`` buffer, + so its per-step windows are intentionally non-contiguous. This kernel indexes + ``(dim, win)`` elements through the view's strides instead of flat-copying. + ``dst`` (the real conv-state pool) is the usual contiguous + ``[layers, cache, dim, K-1]``. + """ + total_requests = step_indices_raw.shape[0] + if total_requests == 0: + return + + if not (dst.is_cuda and src.is_cuda and dst.device == src.device): + raise ValueError( + "fused_conv_window_scatter_with_mask requires dst and src to be CUDA " + f"tensors on the same device ({dst.device=}, {src.device=})." + ) + if dst.ndim != 4 or src.ndim != 5: + raise ValueError(f"Unexpected ranks: {dst.ndim=} (want 4) {src.ndim=} (want 5)") + if dst.shape[0] != src.shape[0]: + raise ValueError(f"Layer dim mismatch: {dst.shape[0]=} vs {src.shape[0]=}") + if dst.shape[2:] != src.shape[3:]: + raise ValueError(f"Window dims mismatch: {dst.shape[2:]=} vs {src.shape[3:]=}") + if dst_indices_raw.ndim != 1 or step_indices_raw.ndim != 1: + raise ValueError( + f"indices must be 1D: {dst_indices_raw.shape=} {step_indices_raw.shape=}" + ) + if dst_indices_raw.shape[0] != step_indices_raw.shape[0]: + raise ValueError( + f"indices length mismatch: {dst_indices_raw.shape[0]=} vs {step_indices_raw.shape[0]=}" + ) + + num_layers = dst.shape[0] + dim = dst.shape[2] + km1 = dst.shape[3] + elem_per_entry = dim * km1 + + src_req_size = src.shape[1] + src_step_size = src.shape[2] + dst_req_size = dst.shape[1] + + # `dst` stays contiguous; `src` is an intentionally non-contiguous (overlapping) + # view, so we do NOT assert src contiguity here (unlike the dense scatter). + if not dst.is_contiguous(): + raise ValueError( + "dst tensor in fused_conv_window_scatter_with_mask must be contiguous" + ) + + dst_indices_raw = dst_indices_raw.to(torch.int32).contiguous() + step_indices_raw = step_indices_raw.to(torch.int32).contiguous() + + BLOCK_SIZE = 1024 + grid = (total_requests, num_layers, triton.cdiv(elem_per_entry, BLOCK_SIZE)) + + _fused_conv_window_scatter_with_mask_kernel[grid]( + src, + dst, + dst_indices_raw, + step_indices_raw, + elem_per_entry, + km1, + src.stride(0), + src.stride(1), + src.stride(2), + src.stride(3), + src.stride(4), + dst.stride(0), + dst.stride(1), + src_req_size, + src_step_size, + dst_req_size, + BLOCK_SIZE=BLOCK_SIZE, + ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index ec8479e1f..b9f549e46 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -90,6 +90,26 @@ _is_fp8_fnuz = is_fp8_fnuz() _use_aiter = bool(envs.SGLANG_USE_AITER.get()) and _is_hip +def conv_window_dedup_enabled( + is_npu: bool, is_cpu: bool, speculative_eagle_topk: Optional[int] +) -> bool: + """Whether the deduplicated sliding-window conv-intermediate layout is safe. + + It is only correct for a *linear* draft chain (``speculative_eagle_topk <= 1``, + i.e. NEXTN / MTP): consecutive draft tokens then form a true sliding window, so + the overlapping physical columns hold identical values. Under EAGLE *tree* + verify (``topk > 1``) the conv kernel walks per-token tree ancestors, so aliased + columns can need different values from different parent chains -> fall back to + the dense layout. NPU/CPU also keep the dense layout (their kernels assume + contiguous per-step windows). See ``MambaPool.__init__``. + """ + return ( + not is_npu + and not is_cpu + and (speculative_eagle_topk is None or speculative_eagle_topk <= 1) + ) + + def get_tensor_size_bytes(t: Union[torch.Tensor, List[torch.Tensor]]): if isinstance(t, list): return sum(get_tensor_size_bytes(x) for x in t) @@ -316,6 +336,7 @@ class MambaPool: device: str, enable_memory_saver: bool = False, speculative_num_draft_tokens: Optional[int] = None, + speculative_eagle_topk: Optional[int] = None, ): conv_state_shape = cache_params.shape.conv temporal_state_shape = cache_params.shape.temporal @@ -393,22 +414,86 @@ class MambaPool: dtype=ssm_dtype, device="cuda", ) - # Cache intermediate conv windows (last K-1 inputs) per draft token during target verify - # Shape: [num_layers, size + 1, speculative_num_draft_tokens, dim, K-1] - intermediate_conv_window_cache = [ - torch.zeros( - size=( - num_mamba_layers, - spec_state_size + 1, - speculative_num_draft_tokens, - conv_shape[0], - conv_shape[1], - ), - dtype=conv_dtype, - device="cuda", - ) - for conv_shape in conv_state_shape - ] + # Cache intermediate conv windows (last K-1 inputs) per draft token + # during target verify. + # + # On CUDA (Triton conv kernel + Triton scatter) we use a + # *deduplicated sliding-window* layout: consecutive draft tokens' + # (K-1)-wide windows overlap by (K-2), so instead of D separate + # [dim, K-1] windows we store one shared [dim, D+K-2] buffer per + # (layer, slot) and expose an overlapping `as_strided` view of + # logical shape [num_layers, size+1, draft_tokens, dim, K-1] where + # step `t`'s window is the slice shared[..., :, t:t+K-1]. This + # halves the conv-intermediate footprint (D*(K-1) -> D+K-2 columns) + # with no numerical change: both the conv kernel write (idempotent + # overlapping stores) and `fused_conv_window_scatter_with_mask` + # consume the view through its strides. + # + # Dedup the sliding-window conv-intermediate only when it is safe: + # CUDA + a linear draft chain (topk <= 1). NPU/CPU and EAGLE tree + # verify (topk > 1) keep the dense layout -- see + # `conv_window_dedup_enabled` for the full rationale. The + # `fused_conv_window_scatter_with_mask` scatter is layout-agnostic, + # so the dense fallback reads correctly through the same code path. + dedup_conv_window = conv_window_dedup_enabled( + _is_npu, _is_cpu, speculative_eagle_topk + ) + self._intermediate_conv_window_phys = [] + if dedup_conv_window: + intermediate_conv_window_cache = [] + for conv_shape in conv_state_shape: + conv_dim, win = conv_shape # win == conv_kernel - 1 == K-1 + shared_win = ( + speculative_num_draft_tokens + win - 1 + ) # D + (K-1) - 1 + phys = torch.zeros( + size=( + num_mamba_layers, + spec_state_size + 1, + conv_dim, + shared_win, + ), + dtype=conv_dtype, + device="cuda", + ) + # view[l, s, step, d, w] = phys[l, s, d, step + w] + view = phys.as_strided( + ( + phys.shape[0], + phys.shape[1], + speculative_num_draft_tokens, + conv_dim, + win, + ), + ( + phys.stride(0), + phys.stride(1), + phys.stride(3), # step -> shared-win axis (stride 1) + phys.stride(2), # dim + phys.stride(3), # win -> shared-win axis (stride 1) + ), + ) + self._intermediate_conv_window_phys.append(phys) + intermediate_conv_window_cache.append(view) + else: + # Original dense layout (NPU/CPU, or EAGLE tree verify): one + # [dim, K-1] window per draft token. + # Shape: [num_layers, size+1, draft_tokens, dim, K-1] + intermediate_conv_window_cache = [ + torch.zeros( + size=( + num_mamba_layers, + spec_state_size + 1, + speculative_num_draft_tokens, + conv_shape[0], + conv_shape[1], + ), + dtype=conv_dtype, + device="cuda", + ) + for conv_shape in conv_state_shape + ] + self._intermediate_conv_window_phys = intermediate_conv_window_cache self.mamba_cache = self.SpeculativeState( conv=conv_state, temporal=temporal_state, @@ -421,7 +506,9 @@ class MambaPool: f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, " f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB " f"intermediate_ssm_state_cache size: {get_tensor_size_bytes(intermediate_ssm_state_cache) / GB:.2f}GB " - f"intermediate_conv_window_cache size: {get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB " + # Report the deduplicated PHYSICAL conv-window buffers (the view + # over-reports its logical, un-deduplicated size). + f"intermediate_conv_window_cache size: {get_tensor_size_bytes(self._intermediate_conv_window_phys) / GB:.2f}GB " ) else: self.mamba_cache = self.State(conv=conv_state, temporal=temporal_state) @@ -431,7 +518,19 @@ class MambaPool: f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, " f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB " ) - self.mem_usage = self.mamba_cache.mem_usage_bytes() / GB + mem_usage_bytes = self.mamba_cache.mem_usage_bytes() + if isinstance(self.mamba_cache, self.SpeculativeState): + # `intermediate_conv_window` is an as_strided view whose logical + # shape over-reports its real footprint; charge the physical buffers + # instead. No-op for the dense layout, where the view and the + # physical tensors coincide. + mem_usage_bytes -= get_tensor_size_bytes( + self.mamba_cache.intermediate_conv_window + ) + mem_usage_bytes += get_tensor_size_bytes( + self._intermediate_conv_window_phys + ) + self.mem_usage = mem_usage_bytes / GB self.num_mamba_layers = num_mamba_layers def get_speculative_mamba2_params_all_layers(self) -> SpeculativeState: @@ -561,6 +660,7 @@ class HybridReqToTokenPool(ReqToTokenPool): enable_mamba_extra_buffer: bool, enable_mamba_extra_buffer_lazy: bool = False, speculative_num_draft_tokens: int = None, + speculative_eagle_topk: Optional[int] = None, enable_overlap_schedule: bool = True, start_layer: Optional[int] = None, ): @@ -585,6 +685,7 @@ class HybridReqToTokenPool(ReqToTokenPool): device=device, enable_mamba_extra_buffer=enable_mamba_extra_buffer, speculative_num_draft_tokens=speculative_num_draft_tokens, + speculative_eagle_topk=speculative_eagle_topk, ) def _init_mamba_pool( @@ -596,6 +697,7 @@ class HybridReqToTokenPool(ReqToTokenPool): device: str, enable_mamba_extra_buffer: bool, speculative_num_draft_tokens: int = None, + speculative_eagle_topk: Optional[int] = None, ): self.mamba_pool = MambaPool( size=mamba_size, @@ -605,6 +707,7 @@ class HybridReqToTokenPool(ReqToTokenPool): device=device, enable_memory_saver=self.enable_memory_saver, speculative_num_draft_tokens=speculative_num_draft_tokens, + speculative_eagle_topk=speculative_eagle_topk, ) self.mamba_allocator = MambaSlotAllocator( size=mamba_size, diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index a86fd9a2c..1717ac95e 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -350,6 +350,7 @@ class ModelRunnerKVCacheMixin: ] ), speculative_num_draft_tokens=max_spec_draft_tokens, + speculative_eagle_topk=self.server_args.speculative_eagle_topk, enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), pre_alloc_size=pre_alloc_size, enable_overlap_schedule=not self.server_args.disable_overlap_schedule, @@ -385,6 +386,7 @@ class ModelRunnerKVCacheMixin: enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), enable_mamba_extra_buffer_lazy=self.server_args.enable_mamba_extra_buffer_lazy(), speculative_num_draft_tokens=max_spec_draft_tokens, + speculative_eagle_topk=self.server_args.speculative_eagle_topk, enable_overlap_schedule=not self.server_args.disable_overlap_schedule, start_layer=self.start_layer, )