From 1ba7c79761534ca56d34b6c7389ccd141405e4e7 Mon Sep 17 00:00:00 2001 From: Ma Mingfei Date: Fri, 26 Jun 2026 07:51:07 +0800 Subject: [PATCH] [CPU] add indices in chunk_gated_delta_rule (#29267) --- .../layers/attention/linear/gdn_backend.py | 2 +- .../attention/linear/kernels/gdn_triton.py | 3 ++- .../srt/model_executor/cpu_graph_runner.py | 3 ++- python/sglang/srt/models/qwen3_5.py | 2 +- sgl-kernel/csrc/cpu/mamba/fla.cpp | 27 ++++++++++++++----- sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 5 ++-- sgl-kernel/python/sgl_kernel/mamba.py | 2 ++ test/registered/cpu/test_mamba.py | 22 +++++++++++---- 8 files changed, 49 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index d97c09880..36995f98a 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -503,7 +503,7 @@ class GDNAttnBackend(MambaAttnBackendBase): query_start_loc=query_start_loc, ) - if (is_npu() or is_cpu()) and last_recurrent_state is not None: + if is_npu() and last_recurrent_state is not None: last_recurrent_state = last_recurrent_state.to( ssm_states.dtype, copy=False ) diff --git a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py index 0daf8f4e2..d037d6e00 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py @@ -141,9 +141,10 @@ class TritonGDNKernel(LinearAttnKernelBase): ) -> tuple: recurrent_state = ssm_states recurrent_state_indices_args = {"initial_state_indices": cache_indices} - if is_npu() or is_cpu(): + if is_npu(): recurrent_state = ssm_states[cache_indices] recurrent_state_indices_args = {} + return chunk_gated_delta_rule( q=q, k=k, diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 5c265f3a9..23f2e9596 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -534,7 +534,8 @@ def register_fake_ops(tp_size: int): cu_seqlens, head_first, use_qk_l2norm_in_kernel, - eps, + initial_state_indices, + eps=1e-6, ): output = torch.empty_like(value) assert initial_state is not None diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 8b1f8052e..dcab16b88 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -1023,7 +1023,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): attn_output = self.attn(q, k, v, forward_batch) if self.attn_output_gate: - if not _is_npu: + if not (_is_npu or _is_cpu): attn_output = fused_sigmoid_mul(attn_output, gate, inplace=True) else: gate_val = gate.reshape(gate.shape[0], -1) if gate.ndim == 3 else gate diff --git a/sgl-kernel/csrc/cpu/mamba/fla.cpp b/sgl-kernel/csrc/cpu/mamba/fla.cpp index a05fb1a15..de6064179 100644 --- a/sgl-kernel/csrc/cpu/mamba/fla.cpp +++ b/sgl-kernel/csrc/cpu/mamba/fla.cpp @@ -864,6 +864,7 @@ template void chunk_gated_delta_rule_fwd_inter_kernel_impl( scalar_t* __restrict__ out, float* __restrict__ state, + const int32_t* __restrict__ indices, const scalar_t* __restrict__ q, const scalar_t* __restrict__ k, const scalar_t* __restrict__ w, @@ -956,7 +957,7 @@ void chunk_gated_delta_rule_fwd_inter_kernel_impl( apply_mask_kernel::apply(attn2, attn, nullptr, d_ptr, mb_size); // step 2.a: v' = w @ state (fuse state *= exp(g_last) with packing) - float* __restrict__ s_ptr = state + bs * (Hv * D * D) + hv * (D * D); + float* __restrict__ s_ptr = state + indices[bs] * (Hv * D * D) + hv * (D * D); const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * (CHUNK_SIZE); float g_last = g_ptr[mb_size - 1]; pack_vnni2( @@ -1475,6 +1476,7 @@ std::tuple chunk_gated_delta_rule_fwd_intra( chunk_gated_delta_rule_fwd_inter_kernel_impl( \ o.data_ptr(), \ initial_state.data_ptr(), \ + initial_state_indices.data_ptr(), \ q.data_ptr(), \ k.data_ptr(), \ w.data_ptr(), \ @@ -1502,14 +1504,15 @@ std::tuple chunk_gated_delta_rule_fwd_inter( const at::Tensor& initial_state, bool output_final_state, const at::Tensor& cu_seqlens, - const at::Tensor& chunk_offsets) { + const at::Tensor& chunk_offsets, + const at::Tensor& initial_state_indices) { const int64_t B = q.size(0); const int64_t T = q.size(1); const int64_t H = q.size(2); const int64_t D = q.size(3); const int64_t Hv = w.size(2); const int64_t Dv = u.size(3); - const int64_t num_seqs = initial_state.size(0); + const int64_t num_seqs = initial_state_indices.size(0); at::Tensor o = at::empty({B, T, Hv, Dv}, q.options()); AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "chunk_gated_delta_rule_fwd_inter", [&] { @@ -1542,6 +1545,7 @@ std::tuple chunk_gated_delta_rule_cpu( const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel, + const at::Tensor& initial_state_indices, double eps = 1e-6) { TORCH_CHECK(!head_first, "chunk_gated_delta_rule_cpu: does not support head first"); @@ -1551,7 +1555,7 @@ std::tuple chunk_gated_delta_rule_cpu( int64_t D = query.size(3); int64_t Hv = value.size(2); int64_t Dv = value.size(3); - int64_t num_seqs = initial_state.size(0); + int64_t num_seqs = initial_state_indices.size(0); TORCH_CHECK(B == 1, __func__, ": expect batch size to be 1"); TORCH_CHECK(Hv % H == 0, __func__, ": expect num_heads_kv multiple of num_heads."); @@ -1564,7 +1568,8 @@ std::tuple chunk_gated_delta_rule_cpu( CHECK_INPUT_SHAPE_DTYPE(g, {B, T, Hv}, at::kFloat); CHECK_INPUT_SHAPE_DTYPE(beta, {B, T, Hv}, at::kBFloat16); CHECK_INPUT_SHAPE_DTYPE(cu_seqlens, {num_seqs + 1}, at::kInt); - CHECK_INPUT_SHAPE_DTYPE(initial_state, {num_seqs, Hv, Dv, D}, at::kFloat); + CHECK_INPUT_SHAPE_DTYPE(initial_state, {initial_state.size(0), Hv, Dv, D}, at::kFloat); + CHECK_INPUT_SHAPE_DTYPE(initial_state_indices, {num_seqs}, at::kInt); constexpr int CHUNK_SIZE = 64; @@ -1582,7 +1587,17 @@ std::tuple chunk_gated_delta_rule_cpu( // fused `chunk_gated_delta_rule_fwd_h` + `chunk_fwd_o` auto [output, final_state] = chunk_gated_delta_rule_fwd_inter( - query_, key_, w, u, g_, decay_mask, initial_state, output_final_state, cu_seqlens, chunk_offsets); + query_, + key_, + w, + u, + g_, + decay_mask, + initial_state, + output_final_state, + cu_seqlens, + chunk_offsets, + initial_state_indices); return std::make_tuple(output, final_state); } diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 3ec8fcdc2..4885bba8d 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -147,7 +147,8 @@ std::tuple chunk_gated_delta_rule_cpu( const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel, - double eps = 1e-5); + const at::Tensor& initial_state_indices, + double eps = 1e-6); // weight prepack at::Tensor convert_weight_packed(at::Tensor& weight); @@ -525,7 +526,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def( "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, Tensor g, Tensor beta, " "Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool head_first, " - "bool use_qk_l2norm_in_kernel, float eps=1e-5) -> (Tensor, Tensor)"); + "bool use_qk_l2norm_in_kernel, Tensor initial_state_indices, float eps=1e-6) -> (Tensor, Tensor)"); m.impl("chunk_gated_delta_rule_cpu", torch::kCPU, &chunk_gated_delta_rule_cpu); // weight prepack diff --git a/sgl-kernel/python/sgl_kernel/mamba.py b/sgl-kernel/python/sgl_kernel/mamba.py index a9ffbfcb5..95e37bd65 100644 --- a/sgl-kernel/python/sgl_kernel/mamba.py +++ b/sgl-kernel/python/sgl_kernel/mamba.py @@ -101,6 +101,7 @@ def chunk_gated_delta_rule_cpu( cu_seqlens, head_first, use_qk_l2norm_in_kernel, + initial_state_indices, ): core_attn_out, last_recurrent_state = ( torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu( @@ -114,6 +115,7 @@ def chunk_gated_delta_rule_cpu( cu_seqlens, head_first, use_qk_l2norm_in_kernel, + initial_state_indices, ) ) h = None # Todo: add return h support diff --git a/test/registered/cpu/test_mamba.py b/test/registered/cpu/test_mamba.py index 8e30afd64..f99009016 100644 --- a/test/registered/cpu/test_mamba.py +++ b/test/registered/cpu/test_mamba.py @@ -252,7 +252,7 @@ def torch_gdn_gating(A_log, a, b, dt_bias): class TestMambaAttention(CustomTestCase): def test_chunk_gated_delta_rule(self): - B, T_PER_SEQ, HK, HV, K, V, N = 1, 128, 16, 32, 128, 128, 4 + B, T_PER_SEQ, HK, HV, K, V, POOL_SIZE = 1, 128, 16, 32, 128, 128, 17 seq_lens = torch.tensor( [T_PER_SEQ - 7, T_PER_SEQ + 11, T_PER_SEQ - 13, T_PER_SEQ + 9], dtype=torch.int32, @@ -264,12 +264,14 @@ class TestMambaAttention(CustomTestCase): ] ) T = cu_seqlens_[-1].item() + cache_indices = torch.tensor([3, 11, 15, 7], dtype=torch.int32) + state_slots = cache_indices query_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16) key_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16) value_ = torch.randn((B, T, HV, V), dtype=torch.bfloat16) g_ = F.logsigmoid(torch.randn((B, T, HV), dtype=torch.float32)) beta_ = torch.sigmoid(torch.randn((B, T, HV), dtype=torch.bfloat16)) - initial_state_ = torch.randn((N, HV, V, K), dtype=torch.float32) * 0.1 + initial_state_ = torch.randn((POOL_SIZE, HV, V, K), dtype=torch.float32) * 0.1 # skip `use_qk_l2norm_in_kernel=False` case since it's not numerically stable in bfloat16 for use_qk_l2norm_in_kernel in [True]: @@ -280,7 +282,7 @@ class TestMambaAttention(CustomTestCase): g=g_, beta=beta_, cu_seqlens=cu_seqlens_, - initial_state=initial_state_, + initial_state=initial_state_[state_slots], use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) @@ -291,8 +293,9 @@ class TestMambaAttention(CustomTestCase): beta = beta_.clone() cu_seqlens = cu_seqlens_.clone() initial_state = initial_state_.clone().transpose(-1, -2).contiguous() + initial_state_before = initial_state.clone() - core_attn_out, last_recurrent_state = ( + core_attn_out, returned_state = ( torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu( query=query, key=key, @@ -304,9 +307,14 @@ class TestMambaAttention(CustomTestCase): cu_seqlens=cu_seqlens, head_first=False, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + initial_state_indices=cache_indices, ) ) - last_recurrent_state = last_recurrent_state.transpose(-1, -2).contiguous() + last_recurrent_state = ( + initial_state[state_slots].transpose(-1, -2).contiguous() + ) + untouched_slots = torch.ones(POOL_SIZE, dtype=torch.bool) + untouched_slots[state_slots] = False atol = rtol = precision[core_attn_out.dtype] torch.testing.assert_close( core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol @@ -314,6 +322,10 @@ class TestMambaAttention(CustomTestCase): torch.testing.assert_close( last_recurrent_state, last_recurrent_state_ref, atol=atol, rtol=rtol ) + torch.testing.assert_close(returned_state, initial_state) + torch.testing.assert_close( + initial_state[untouched_slots], initial_state_before[untouched_slots] + ) def test_fused_gdn_gating(self): dims = [6, 32]