minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[CPU] add indices in chunk_gated_delta_rule (#29267)

Browse Source
This commit is contained in:
Ma Mingfei
2026-06-26 07:51:07 +08:00
committed by GitHub
parent c8c6757bed
commit 1ba7c79761
8 changed files with 49 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/linear/gdn_backend.py
+1 -1
View File
@@ -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
+2 -1
View File
@@ -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
+2 -1
View File
@@ -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
+1 -1
View File
@@ -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
+21 -6
View File
@@ -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
+3 -2
View File
@@ -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
+2
View File
@@ -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
+17 -5
View File
@@ -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]