[CPU] add indices in chunk_gated_delta_rule (#29267)

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
                 )

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,

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

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

sgl-kernel/csrc/cpu/mamba/fla.cpp

@@ -864,6 +864,7 @@ template <typename scalar_t, int D, int CHUNK_SIZE>
 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<scalar_t, CHUNK_SIZE, false>::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<scalar_t, D, D>(
@@ -1475,6 +1476,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_intra(
   chunk_gated_delta_rule_fwd_inter_kernel_impl<scalar_t, HD, CHUNK_SIZE>( \
       o.data_ptr<scalar_t>(),                                             \
       initial_state.data_ptr<float>(),                                    \
+      initial_state_indices.data_ptr<int32_t>(),                          \
       q.data_ptr<scalar_t>(),                                             \
       k.data_ptr<scalar_t>(),                                             \
       w.data_ptr<scalar_t>(),                                             \
@@ -1502,14 +1504,15 @@ std::tuple<at::Tensor, at::Tensor> 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<at::Tensor, at::Tensor> 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<at::Tensor, at::Tensor> 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<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
   CHECK_INPUT_SHAPE_DTYPE<false>(g, {B, T, Hv}, at::kFloat);
   CHECK_INPUT_SHAPE_DTYPE<false>(beta, {B, T, Hv}, at::kBFloat16);
   CHECK_INPUT_SHAPE_DTYPE<false>(cu_seqlens, {num_seqs + 1}, at::kInt);
-  CHECK_INPUT_SHAPE_DTYPE<false>(initial_state, {num_seqs, Hv, Dv, D}, at::kFloat);
+  CHECK_INPUT_SHAPE_DTYPE<false>(initial_state, {initial_state.size(0), Hv, Dv, D}, at::kFloat);
+  CHECK_INPUT_SHAPE_DTYPE<false>(initial_state_indices, {num_seqs}, at::kInt);
 
   constexpr int CHUNK_SIZE = 64;
 
@@ -1582,7 +1587,17 @@ std::tuple<at::Tensor, at::Tensor> 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<CHUNK_SIZE>(
-      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);
 }

sgl-kernel/csrc/cpu/torch_extension_cpu.cpp

@@ -147,7 +147,8 @@ std::tuple<at::Tensor, at::Tensor> 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

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

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]