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

[CPU] update fla.cpp to support when num_head_v is not multiples of 16 (#30604)

Browse Source
This commit is contained in:
Ma Mingfei
2026-07-10 09:21:07 +08:00
committed by GitHub
parent 5ce5e1ee3e
commit 073b36853f
2 changed files with 218 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/mamba/fla.cpp
+25 -9
View File
@@ -135,8 +135,13 @@ struct l2norm_kernel<at::BFloat16, D, has_scale> {
template <typename scalar_t, int CHUNK_SIZE, int BLOCK_H>
struct cumsum_kernel {
static inline void
apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int size, int ld_src, int ld_dst) {
static inline void apply(
scalar_t* __restrict__ out,
const scalar_t* __restrict__ input,
int mb_size,
int hb_size,
int ld_src,
int ld_dst) {
TORCH_CHECK(false, "cumsum_kernel: scalar path not implemented!");
}
};
@@ -144,9 +149,12 @@ struct cumsum_kernel {
#if defined(CPU_CAPABILITY_AVX512)
template <int CHUNK_SIZE, int BLOCK_H>
struct cumsum_kernel<float, CHUNK_SIZE, BLOCK_H> {
static inline void apply(float* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) {
static inline void
apply(float* __restrict__ out, const float* __restrict__ input, int mb_size, int hb_size, int ld_src, int ld_dst) {
// vector length of fp32 for avx512
static_assert(BLOCK_H == 16);
TORCH_CHECK(hb_size > 0 && hb_size <= BLOCK_H);
const __mmask16 vmask = static_cast<__mmask16>((1u << hb_size) - 1u);
__m512i va[16];
__m512 vsum = _mm512_set1_ps(0.f);
@@ -154,14 +162,23 @@ struct cumsum_kernel<float, CHUNK_SIZE, BLOCK_H> {
for (int i = 0; i < CHUNK_SIZE; i += 16) {
// load input data
Unroll<16>{}([&](auto j) {
__m512 v = (i + j < size) ? _mm512_loadu_ps(input + (i + j) * ld_src) : _mm512_setzero_ps();
__m512 v;
if (i + j < mb_size) {
v = _mm512_maskz_loadu_ps(vmask, input + (i + j) * ld_src);
} else {
v = _mm512_setzero_ps();
}
vsum = _mm512_add_ps(vsum, v);
va[j] = _mm512_castps_si512(vsum);
});
// transpose
transpose_16x16_32bit(va);
// store output data
Unroll<16>{}([&](auto j) { _mm512_storeu_si512(out + j * ld_dst + i, va[j]); });
Unroll<16>{}([&](auto j) {
if (j < hb_size) {
_mm512_storeu_si512(out + j * ld_dst + i, va[j]);
}
});
}
}
};
@@ -633,9 +650,7 @@ void chunk_local_cumsum_kernel_impl(
int64_t Hv,
int64_t NT) {
constexpr int BLOCK_H = 16;
// TODO: now we only support qwen3.5 configs (H/Hv == 16/32)
TORCH_CHECK(Hv % BLOCK_H == 0);
int64_t HB = Hv / BLOCK_H;
int64_t HB = div_up(Hv, int64_t(BLOCK_H));
// parallel on [NT * HB] to increase parallelism
at::parallel_for(0, NT * HB, 0, [&](int64_t begin, int64_t end) {
@@ -648,10 +663,11 @@ void chunk_local_cumsum_kernel_impl(
int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs];
int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE;
int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE));
int64_t hb_size = std::min(Hv - hb * BLOCK_H, int64_t(BLOCK_H));
const scalar_t* __restrict__ g_ptr = g + (batch_offset + mb_start) * Hv + hb * BLOCK_H;
scalar_t* __restrict__ gsum_ptr = g_ + nt * (Hv * CHUNK_SIZE) + hb * (BLOCK_H * CHUNK_SIZE);
cumsum_kernel<scalar_t, CHUNK_SIZE, BLOCK_H>::apply(gsum_ptr, g_ptr, mb_size, Hv, CHUNK_SIZE);
cumsum_kernel<scalar_t, CHUNK_SIZE, BLOCK_H>::apply(gsum_ptr, g_ptr, mb_size, hb_size, Hv, CHUNK_SIZE);
// move to the next index
data_index_step(nt, NT, hb, HB);
test/registered/cpu/test_mamba.py
+45 -44
View File
@@ -1,12 +1,12 @@
import unittest
import sys
import pytest
import torch
import torch.nn.functional as F
from torch.nn.functional import softplus
from utils import parametrize, precision
from utils import precision
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
@@ -250,9 +250,14 @@ def torch_gdn_gating(A_log, a, b, dt_bias):
), b.sigmoid().unsqueeze(0)
class TestMambaAttention(CustomTestCase):
def test_chunk_gated_delta_rule(self):
B, T_PER_SEQ, HK, HV, K, V, POOL_SIZE = 1, 128, 16, 32, 128, 128, 17
@pytest.mark.parametrize(
("B", "T_PER_SEQ", "HK", "HV", "K", "V", "POOL_SIZE"),
[
(1, 128, 3, 6, 128, 128, 17),
(1, 128, 16, 32, 128, 128, 17),
],
)
def test_chunk_gated_delta_rule(B, T_PER_SEQ, HK, HV, K, V, POOL_SIZE):
seq_lens = torch.tensor(
[T_PER_SEQ - 7, T_PER_SEQ + 11, T_PER_SEQ - 13, T_PER_SEQ + 9],
dtype=torch.int32,
@@ -295,8 +300,7 @@ class TestMambaAttention(CustomTestCase):
initial_state = initial_state_.clone().transpose(-1, -2).contiguous()
initial_state_before = initial_state.clone()
core_attn_out, returned_state = (
torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu(
core_attn_out, returned_state = torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu(
query=query,
key=key,
value=value,
@@ -309,10 +313,7 @@ class TestMambaAttention(CustomTestCase):
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
initial_state_indices=cache_indices,
)
)
last_recurrent_state = (
initial_state[state_slots].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]
@@ -327,35 +328,43 @@ class TestMambaAttention(CustomTestCase):
initial_state[untouched_slots], initial_state_before[untouched_slots]
)
def test_fused_gdn_gating(self):
dims = [6, 32]
for dim in dims:
for A_log_dtype in [torch.float32, torch.bfloat16]:
@pytest.mark.parametrize("dim", [6, 32])
@pytest.mark.parametrize(
"A_log_dtype",
[torch.float32, torch.bfloat16],
ids=["float32", "bfloat16"],
)
def test_fused_gdn_gating(dim, A_log_dtype):
A_log = torch.rand(dim, dtype=A_log_dtype)
a = torch.rand(1024, dim, dtype=torch.bfloat16)
b = torch.rand(1024, dim, dtype=torch.bfloat16)
dt_bias = torch.rand(dim, dtype=torch.bfloat16)
g, beta = torch_gdn_gating(A_log, a, b, dt_bias)
g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(
A_log, a, b, dt_bias
)
g_sgl, beta_sgl = torch.ops.sgl_kernel.fused_gdn_gating_cpu(A_log, a, b, dt_bias)
atol = rtol = precision[g.dtype]
atol2 = rtol2 = precision[beta.dtype]
torch.testing.assert_close(g, g_sgl, atol=atol, rtol=rtol)
torch.testing.assert_close(beta, beta_sgl, atol=atol2, rtol=rtol2)
@parametrize(
batch_size=[1, 4],
num_value_heads=[32],
head_k_dim=[128],
head_v_dim=[128],
num_heads=[16],
seq_len=[1],
attn_tp_size=[1],
)
def test_fused_sigmoid_gating_delta_rule_update(
self,
@pytest.mark.parametrize(
(
"batch_size",
"num_value_heads",
"head_k_dim",
"head_v_dim",
"num_heads",
"seq_len",
"attn_tp_size",
),
[
(1, 32, 128, 128, 16, 1, 1),
(4, 32, 128, 128, 16, 1, 1),
],
)
def test_fused_sigmoid_gating_delta_rule_update(
batch_size,
num_value_heads,
head_k_dim,
@@ -363,13 +372,11 @@ class TestMambaAttention(CustomTestCase):
num_heads,
seq_len,
attn_tp_size,
):
):
key_dim = head_k_dim * num_heads
value_dim = head_v_dim * num_value_heads
mixed_qkv_dim = (key_dim * 2 + value_dim) // attn_tp_size
mixed_qkv = torch.rand(
seq_len * batch_size, mixed_qkv_dim, dtype=torch.bfloat16
)
mixed_qkv = torch.rand(seq_len * batch_size, mixed_qkv_dim, dtype=torch.bfloat16)
query, key, value = torch.split(
mixed_qkv,
[
@@ -399,8 +406,7 @@ class TestMambaAttention(CustomTestCase):
key_ref = key_ref.repeat_interleave(num_value_heads // num_heads, dim=2)
for A_log_dtype in [torch.float32, torch.bfloat16]:
A_log = A_log.to(A_log_dtype)
core_attn_out_ref, last_recurrent_state_ref = (
sigmoid_gating_delta_rule_update(
core_attn_out_ref, last_recurrent_state_ref = sigmoid_gating_delta_rule_update(
query_ref.transpose(0, 1),
key_ref.transpose(0, 1),
value.transpose(0, 1),
@@ -408,15 +414,11 @@ class TestMambaAttention(CustomTestCase):
a,
dt_bias,
b,
initial_state=ssm_states_kv[cache_indices]
.transpose(-1, -2)
.contiguous(),
initial_state=ssm_states_kv[cache_indices].transpose(-1, -2).contiguous(),
output_final_state=True,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
)
)
core_attn_out = (
torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
core_attn_out = torch.ops.sgl_kernel.fused_sigmoid_gating_delta_rule_update_cpu(
A_log=A_log,
dt_bias=dt_bias,
q=query,
@@ -431,7 +433,6 @@ class TestMambaAttention(CustomTestCase):
softplus_beta=1.0,
softplus_threshold=20.0,
)
)
last_recurrent_state = (
ssm_states_kv[cache_indices].transpose(-1, -2).contiguous()
)
@@ -445,4 +446,4 @@ class TestMambaAttention(CustomTestCase):
if __name__ == "__main__":
unittest.main()
sys.exit(pytest.main([__file__]))