Fix unnecessary gather/scatter on CPU for non-contiguous Mamba statepool (#31754)

This commit is contained in:
YanbingJiang
2026-07-23 09:08:45 +08:00
committed by GitHub
parent 9bda6fdb9a
commit 9a7ac3ecef
@@ -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()