From 9a7ac3ecefc64bf4b47c70e3446e233ac02773ea Mon Sep 17 00:00:00 2001 From: YanbingJiang Date: Thu, 23 Jul 2026 09:08:45 +0800 Subject: [PATCH] Fix unnecessary gather/scatter on CPU for non-contiguous Mamba statepool (#31754) --- python/sglang/srt/layers/attention/linear/gdn_backend.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 64731469c..d3dd717be 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -488,11 +488,16 @@ class GDNAttnBackend(MambaAttnBackendBase): # slot layout, so they silently drop the write to the strided envelope # pool. Run them on contiguous per-sequence copies (identity-indexed) and # scatter the result back. No-op for the default contiguous pool. + # CPU kernels (causal_conv1d_fwd_cpu, chunk_gated_delta_rule_cpu) use + # proper indexed writes and handle non-contiguous pools directly via + # cache_indices, so the gather/scatter round-trip is unnecessary on CPU. # TODO(ch-wan): drop these .contiguous() copies by making the prefill conv # and chunk_gated_delta_rule kernels honor the pool's real slot stride + # int64 indexing, like packed_decode / causal_conv1d_update already do. - needs_state_gather = (not is_target_verify) and ( - not conv_states.is_contiguous() or not ssm_states.is_contiguous() + needs_state_gather = ( + (not is_target_verify) + and (not is_cpu()) + and (not conv_states.is_contiguous() or not ssm_states.is_contiguous()) ) if needs_state_gather: conv_states_contig = conv_states[cache_indices].contiguous()