iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CPU] Support FP16/BF16 persisted GDN state on AMX (#52191)

Browse Source 
Signed-off-by: Li, Tianmu <tianmu.li@intel.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Tianmu Li
2026-08-31 16:32:20 +08:00
committed by GitHub
co-authored by OpenAI Codex
parent 2a61f060d3
commit c6c33f2b1f
4 changed files with 575 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/cpu/sgl-kernels/fla.cpp
+292 -118
View File
@@ -74,6 +74,20 @@ void pack_vnni2(scalar_t* __restrict__ dst, float* __restrict__ src, const float
#endif
}
template <typename state_t>
inline void store_state_vec2(
state_t* __restrict__ dst,
const at::vec::Vectorized<float>& src0,
const at::vec::Vectorized<float>& src1) {
if constexpr (std::is_same_v<state_t, float>) {
src0.store(dst);
src1.store(dst + at::vec::Vectorized<float>::size());
} else {
auto state_vec = convert_from_float_ext<state_t>(src0, src1);
state_vec.store(dst);
}
}
template <typename scalar_t, int SIZE>
inline void fill_stub(scalar_t* __restrict__ out, float val) {
using Vec = at::vec::Vectorized<scalar_t>;
@@ -1255,10 +1269,10 @@ void chunk_gated_delta_rule_fwd_intra_kernel_impl(
// d : [B, NT, Hv, C, C]
// cu_seqlens : [num_seqs + 1]
// chunk_offsets : [num_seqs + 1]
template <typename scalar_t, int D, int CHUNK_SIZE>
template <typename scalar_t, typename state_t, int D, int CHUNK_SIZE>
void chunk_gated_delta_rule_fwd_inter_kernel_impl(
scalar_t* __restrict__ out,
float* __restrict__ state,
state_t* __restrict__ state,
const int32_t* __restrict__ indices,
const scalar_t* __restrict__ q,
const scalar_t* __restrict__ k,
@@ -1297,6 +1311,7 @@ void chunk_gated_delta_rule_fwd_inter_kernel_impl(
// thread local temp buffer
DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D);
DECL_ZERO_BUF(scalar_t, tmp2, D * D);
DECL_ZERO_BUF(float, state_tmp, D * D);
DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D);
DECL_ZERO_BUF(scalar_t, tmp4, CHUNK_SIZE * D);
DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE);
@@ -1370,7 +1385,27 @@ 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 + indices[bs] * state_strideS + hv * (D * D);
state_t* __restrict__ state_slot =
state + indices[bs] * state_strideS + hv * (D * D);
float* s_ptr;
if constexpr (std::is_same_v<state_t, float>) {
s_ptr = state_slot;
} else {
using bVec = at::vec::Vectorized<state_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
constexpr int fVecSize = fVec::size();
int64_t d0 = 0;
for (; d0 + kVecSize <= D * D; d0 += kVecSize) {
auto [state_vec0, state_vec1] = load_float_vec2(state_slot + d0);
state_vec0.store(state_tmp + d0);
state_vec1.store(state_tmp + d0 + fVecSize);
}
for (; d0 < D * D; ++d0) {
state_tmp[d0] = static_cast<float>(state_slot[d0]);
}
s_ptr = state_tmp;
}
const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * (CHUNK_SIZE);
float g_last = g_ptr[mb_size - 1];
const scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH;
@@ -1547,6 +1582,22 @@ void chunk_gated_delta_rule_fwd_inter_kernel_impl(
s_ptr,
D);
}
if constexpr (!std::is_same_v<state_t, float>) {
using bVec = at::vec::Vectorized<state_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int kVecSize = bVec::size();
constexpr int fVecSize = fVec::size();
int64_t d0 = 0;
for (; d0 + kVecSize <= D * D; d0 += kVecSize) {
fVec state_vec0 = fVec::loadu(s_ptr + d0);
fVec state_vec1 = fVec::loadu(s_ptr + d0 + fVecSize);
store_state_vec2<state_t>(state_slot + d0, state_vec0, state_vec1);
}
for (; d0 < D * D; ++d0) {
state_slot[d0] = static_cast<state_t>(s_ptr[d0]);
}
}
}
// move to the next index
@@ -1576,7 +1627,7 @@ inline at::vec::Vectorized<float> softplus(const at::vec::Vectorized<float>& x,
return Vec::blendv(Vec::blendv(log1pex, expx, mask_lo), x, mask_hi);
}
template <typename scalar_t, typename param_t>
template <typename scalar_t, typename param_t, typename state_t>
void fused_sigmoid_gating_delta_rule_update_kernel_impl(
const scalar_t* __restrict__ q_ptr,
const scalar_t* __restrict__ k_ptr,
@@ -1586,7 +1637,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
const scalar_t* __restrict__ dt_bias_ptr,
const scalar_t* __restrict__ b_ptr,
const int32_t* __restrict__ indices_ptr,
float* __restrict__ state_ptr,
state_t* __restrict__ state_ptr,
scalar_t* __restrict__ o_ptr,
float* __restrict__ qk_scale_buf,
int64_t seq_len,
@@ -1700,8 +1751,10 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
state_vec1 = state_vec1 * g_val_exp_vec + k_vec * dt_vec1;
o_vec0 = o_vec0 + state_vec0 * q_vec * scale_vec;
o_vec1 = o_vec1 + state_vec1 * q_vec * scale_vec;
state_vec0.store(state_ptr + state_offset + di * v_head_dim + dvi);
state_vec1.store(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
store_state_vec2<state_t>(
state_ptr + state_offset + di * v_head_dim + dvi,
state_vec0,
state_vec1);
}
bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1);
o_vec.store(o_ptr + o_offset + dvi);
@@ -1710,8 +1763,10 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
float kv_mem_val = 0;
for (int di = 0; di < head_dim; ++di) {
float k_val = k_ptr[k_offset + di] * k_scale;
state_ptr[state_offset + di * v_head_dim + dvi] *= g_val_exp;
kv_mem_val += state_ptr[state_offset + di * v_head_dim + dvi] * k_val;
float state_val =
static_cast<float>(state_ptr[state_offset + di * v_head_dim + dvi]) *
g_val_exp;
kv_mem_val += state_val * k_val;
}
float v_val = v_ptr[v_offset + dvi];
float dt_val = (v_val - kv_mem_val) * beta_val;
@@ -1719,8 +1774,13 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
for (int di = 0; di < head_dim; ++di) {
float q_val = q_ptr[q_offset + di] * q_scale;
float k_val = k_ptr[k_offset + di] * k_scale;
state_ptr[state_offset + di * v_head_dim + dvi] += k_val * dt_val;
o_val += state_ptr[state_offset + di * v_head_dim + dvi] * q_val * scale;
float state_val =
static_cast<float>(state_ptr[state_offset + di * v_head_dim + dvi]) *
g_val_exp;
state_val += k_val * dt_val;
state_ptr[state_offset + di * v_head_dim + dvi] =
static_cast<state_t>(state_val);
o_val += state_val * q_val * scale;
}
o_ptr[o_offset + dvi] = o_val;
}
@@ -1736,7 +1796,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
// token ``t`` into cache slot ``t`` (multi-slot rollback, matching the GPU
// kernel). Parallelized over (sequence, v_head); the per-sequence token loop is
// sequential as required by the recurrence.
template <typename scalar_t, typename param_t>
template <typename scalar_t, typename param_t, typename state_t>
void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
const scalar_t* __restrict__ q_ptr, // [T, HK, EK]
const scalar_t* __restrict__ k_ptr, // [T, HK, EK]
@@ -1748,7 +1808,7 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
const int32_t* __restrict__ spec_indices_ptr, // [N, S]
const int32_t* __restrict__ num_accepted_ptr, // [N]
const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1]
float* __restrict__ state_ptr,
state_t* __restrict__ state_ptr,
scalar_t* __restrict__ o_ptr, // [T, HV, EV]
float* __restrict__ qk_scale_buf, // [2, T, HK]
int64_t total_tokens,
@@ -1770,7 +1830,6 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
using bVec = at::vec::Vectorized<scalar_t>;
using fVec = at::vec::Vectorized<float>;
constexpr int64_t VecSize = bVec::size();
constexpr int64_t fVecSize = fVec::size();
int64_t group_size = v_num_heads / num_heads;
double scale = 1 / std::sqrt((double)head_dim);
fVec scale_vec = fVec((float)scale);
@@ -1815,8 +1874,10 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
for (int64_t t = 0; t < q_len; ++t) {
int64_t cur_slot = (int64_t)spec_indices_ptr[bi * spec_stride + t];
int64_t token = q_start + t;
const float* src = state_ptr + prev_slot * state_slot_stride + ni * head_dim * v_head_dim;
float* dst = state_ptr + cur_slot * state_slot_stride + ni * head_dim * v_head_dim;
const state_t* src =
state_ptr + prev_slot * state_slot_stride + ni * head_dim * v_head_dim;
state_t* dst =
state_ptr + cur_slot * state_slot_stride + ni * head_dim * v_head_dim;
float g_val = -std::exp((float)A_log_ptr[ni]) *
softplus((float)a_ptr[token * v_num_heads + ni] + (float)dt_bias_ptr[ni], softplus_threshold);
float g_val_exp = std::exp(g_val);
@@ -1836,8 +1897,8 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
fVec kv_mem_vec1 = fVec(0.f);
for (int di = 0; di < head_dim; ++di) {
fVec k_val_vec = fVec((float)k_ptr[k_offset + di] * k_scale);
fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi);
fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize);
auto [sv0, sv1] =
load_float_vec2(src + di * v_head_dim + dvi);
kv_mem_vec0 = kv_mem_vec0 + sv0 * g_val_exp_vec * k_val_vec;
kv_mem_vec1 = kv_mem_vec1 + sv1 * g_val_exp_vec * k_val_vec;
}
@@ -1851,14 +1912,14 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
for (int di = 0; di < head_dim; ++di) {
fVec q_vec = fVec((float)q_ptr[q_offset + di] * q_scale);
fVec k_vec = fVec((float)k_ptr[k_offset + di] * k_scale);
fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi);
fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize);
auto [sv0, sv1] =
load_float_vec2(src + di * v_head_dim + dvi);
sv0 = sv0 * g_val_exp_vec + k_vec * dt_vec0;
sv1 = sv1 * g_val_exp_vec + k_vec * dt_vec1;
o_vec0 = o_vec0 + sv0 * q_vec * scale_vec;
o_vec1 = o_vec1 + sv1 * q_vec * scale_vec;
sv0.store(dst + di * v_head_dim + dvi);
sv1.store(dst + di * v_head_dim + dvi + fVecSize);
store_state_vec2<state_t>(
dst + di * v_head_dim + dvi, sv0, sv1);
}
bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1);
o_vec.store(o_ptr + o_offset + dvi);
@@ -1867,7 +1928,8 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
float kv_mem_val = 0.f;
for (int di = 0; di < head_dim; ++di) {
float k_val = (float)k_ptr[k_offset + di] * k_scale;
kv_mem_val += src[di * v_head_dim + dvi] * g_val_exp * k_val;
kv_mem_val += static_cast<float>(src[di * v_head_dim + dvi]) *
g_val_exp * k_val;
}
float v_val = (float)v_ptr[v_offset + dvi];
float dt_val = (v_val - kv_mem_val) * beta_val;
@@ -1875,8 +1937,10 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
for (int di = 0; di < head_dim; ++di) {
float q_val = (float)q_ptr[q_offset + di] * q_scale;
float k_val = (float)k_ptr[k_offset + di] * k_scale;
float ns = src[di * v_head_dim + dvi] * g_val_exp + k_val * dt_val;
dst[di * v_head_dim + dvi] = ns;
float ns = static_cast<float>(src[di * v_head_dim + dvi]) *
g_val_exp +
k_val * dt_val;
dst[di * v_head_dim + dvi] = static_cast<state_t>(ns);
o_val += ns * q_val * scale;
}
o_ptr[o_offset + dvi] = (scalar_t)o_val;
@@ -2112,10 +2176,10 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_intra(
return std::make_tuple(w, u, decay_mask);
}
#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \
chunk_gated_delta_rule_fwd_inter_kernel_impl<scalar_t, HD, CHUNK_SIZE>( \
#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \
chunk_gated_delta_rule_fwd_inter_kernel_impl<scalar_t, state_t, HD, CHUNK_SIZE>( \
o.data_ptr<scalar_t>(), \
initial_state.data_ptr<float>(), \
initial_state.data_ptr<state_t>(), \
initial_state_indices.data_ptr<int32_t>(), \
q.data_ptr<scalar_t>(), \
k.data_ptr<scalar_t>(), \
@@ -2134,7 +2198,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_intra(
k.stride(2), \
initial_state.stride(0));
template <int CHUNK_SIZE>
template <typename state_t, int CHUNK_SIZE>
std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_inter(
const at::Tensor& q,
const at::Tensor& k,
@@ -2172,7 +2236,7 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_inter(
// value: [B, T, Hv, Dv]
// g: [B, T, Hv] FP32
// beta: [B, T, Hv]
// initial_state: [num_seqs, Hv, Dv, D] FP32
// initial_state: [num_seqs, Hv, Dv, D] FP16, BF16, or FP32 persisted state
// cu_seqlens: [num_seqs + 1] INT32
//
std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
@@ -2211,7 +2275,12 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
CHECK_INPUT_SHAPE_DTYPE<false>(cu_seqlens, {num_seqs + 1}, at::kInt);
TORCH_CHECK(initial_state.sizes() == at::IntArrayRef({initial_state.size(0), Hv, Dv, D}),
"chunk_gated_delta_rule_cpu: initial_state shape mismatch, got ", initial_state.sizes());
TORCH_CHECK(initial_state.scalar_type() == at::kFloat, "chunk_gated_delta_rule_cpu: initial_state dtype mismatch");
TORCH_CHECK(
initial_state.scalar_type() == at::kFloat ||
initial_state.scalar_type() == at::kBFloat16 ||
initial_state.scalar_type() == at::kHalf,
"chunk_gated_delta_rule_cpu: initial_state dtype must be float16, "
"bfloat16, or float32");
CHECK_CPU(initial_state);
// initial_state may be a pooled/paged buffer with padding between slots
// (e.g. mamba cache-align mode), so only the per-slot (Hv, Dv, D) layout
@@ -2237,18 +2306,28 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
chunk_gated_delta_rule_fwd_intra<CHUNK_SIZE>(key_, value, g_, beta, cu_seqlens, chunk_indices);
// 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,
initial_state_indices);
auto launch_inter = [&](auto state_tag) {
using state_t = decltype(state_tag);
return chunk_gated_delta_rule_fwd_inter<state_t, CHUNK_SIZE>(
query_,
key_,
w,
u,
g_,
decay_mask,
initial_state,
output_final_state,
cu_seqlens,
chunk_offsets,
initial_state_indices);
};
auto output_and_state =
initial_state.scalar_type() == at::kFloat
? launch_inter(float{})
: initial_state.scalar_type() == at::kBFloat16
? launch_inter(at::BFloat16{})
: launch_inter(at::Half{});
auto [output, final_state] = output_and_state;
return std::make_tuple(output, final_state);
}
@@ -2261,6 +2340,7 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
// a: [batch_size, v_num_heads]
// b: [batch_size, v_num_heads]
// initial_state_source:[num_tokens, v_num_heads, head_dim, v_head_dim]
// FP16, BF16, or FP32 persisted state; recurrent arithmetic remains FP32.
// initial_state_indices: [batch_size]
// cu_seqlens: [batch_size + 1]
at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
@@ -2293,8 +2373,24 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
CHECK_INPUT_SHAPE_DTYPE<true>(b, {batch_size, v_num_heads}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(initial_state_indices, {batch_size}, at::kInt);
CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt);
CHECK_INPUT_SHAPE_DTYPE<true>(
initial_state_source, {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}, at::kFloat);
TORCH_CHECK(
initial_state_source.scalar_type() == at::kFloat ||
initial_state_source.scalar_type() == at::kBFloat16 ||
initial_state_source.scalar_type() == at::kHalf,
"fused_sigmoid_gating_delta_rule_update_cpu: initial_state dtype must "
"be float16, bfloat16, or float32");
TORCH_CHECK(
initial_state_source.sizes() ==
at::IntArrayRef(
{initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}),
"Input tensor shape mismatch for initial_state_source");
CHECK_CPU(initial_state_source);
TORCH_CHECK(
initial_state_source.stride(-1) == 1 &&
initial_state_source.stride(-2) == v_head_dim &&
initial_state_source.stride(-3) == head_dim * v_head_dim,
"fused_sigmoid_gating_delta_rule_update_cpu: expect initial_state_source "
"to be contiguous per pool slot.");
CHECK(initial_state_source.size(0) >= batch_size);
CHECK_EQ(v_num_heads % num_heads, 0);
TORCH_CHECK(
@@ -2318,46 +2414,63 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
at::Tensor core_attn_out = at::empty({batch_size, seq_len, v_num_heads, v_head_dim}, q.options());
at::Tensor qk_scale_buf = at::empty({2 * batch_size, seq_len, num_heads}, at::kFloat);
CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] {
fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t, param_t>(
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
A_log.data_ptr<param_t>(),
a.data_ptr<scalar_t>(),
dt_bias.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
initial_state_indices.data_ptr<int32_t>(),
initial_state_source.data_ptr<float>(),
core_attn_out.data_ptr<scalar_t>(),
qk_scale_buf.data_ptr<float>(),
seq_len,
batch_size,
num_heads,
head_dim,
v_num_heads,
v_head_dim,
q_strideB,
q_strideS,
q_strideH,
k_strideB,
k_strideS,
k_strideH,
v_strideB,
v_strideS,
v_strideH,
state_slot_stride,
use_qk_l2norm_in_kernel,
softplus_threshold);
});
auto launch = [&](auto state_tag) {
using state_t = decltype(state_tag);
CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
q.scalar_type(), A_log.scalar_type(),
"fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] {
fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t, param_t, state_t>(
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
A_log.data_ptr<param_t>(),
a.data_ptr<scalar_t>(),
dt_bias.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
initial_state_indices.data_ptr<int32_t>(),
initial_state_source.data_ptr<state_t>(),
core_attn_out.data_ptr<scalar_t>(),
qk_scale_buf.data_ptr<float>(),
seq_len,
batch_size,
num_heads,
head_dim,
v_num_heads,
v_head_dim,
q_strideB,
q_strideS,
q_strideH,
k_strideB,
k_strideS,
k_strideH,
v_strideB,
v_strideS,
v_strideH,
state_slot_stride,
use_qk_l2norm_in_kernel,
softplus_threshold);
});
};
if (initial_state_source.scalar_type() == at::kFloat) {
launch(float());
} else if (initial_state_source.scalar_type() == at::kBFloat16) {
launch(at::BFloat16());
} else if (initial_state_source.scalar_type() == at::kHalf) {
launch(at::Half());
} else {
TORCH_CHECK(
false,
"fused_sigmoid_gating_delta_rule_update_cpu: initial_state dtype must "
"be float16, bfloat16, or float32");
}
return core_attn_out;
}
// Speculative-decode update (multi-token, multi-slot rollback).
// q: [T, HK, EK] k: [T, HK, EK] v: [T, HV, EV]
// a: [T, HV] b: [T, HV]
// initial_state_source: [N_slots, HV, EK, EV] FP32 (updated in place)
// initial_state_source: [N_slots, HV, EK, EV] FP16, BF16, or FP32 persisted state
// (updated in place; recurrent arithmetic remains FP32)
// spec_state_indices: [batch, S] INT32 (S = num_spec + 1)
// num_accepted_tokens: [batch] INT32
// cu_seqlens: [batch + 1] INT32
@@ -2385,21 +2498,65 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(
int64_t head_dim = q.size(2);
int64_t v_num_heads = v.size(1);
int64_t v_head_dim = v.size(2);
int64_t batch_size = cu_seqlens.size(0) - 1;
int64_t spec_stride = spec_state_indices.stride(0);
CHECK_INPUT_SHAPE_DTYPE<true>(k, {total_tokens, num_heads, head_dim}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(v, {total_tokens, v_num_heads, v_head_dim}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(a, {total_tokens, v_num_heads}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(b, {total_tokens, v_num_heads}, q.scalar_type());
CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type());
CHECK_DIM(1, cu_seqlens);
int64_t batch_size = cu_seqlens.size(0) - 1;
CHECK_INPUT_SHAPE_DTYPE<true>(num_accepted_tokens, {batch_size}, at::kInt);
CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt);
CHECK_DIM(2, spec_state_indices);
CHECK_CPU(spec_state_indices);
TORCH_CHECK(
spec_state_indices.scalar_type() == at::kInt,
"spec_state_indices must be an int32 CPU tensor");
TORCH_CHECK(
spec_state_indices.size(0) == batch_size,
"spec_state_indices first dimension must equal batch_size");
CHECK_LAST_DIM_CONTIGUOUS_INPUT(spec_state_indices);
const int32_t* cu_seqlens_ptr = cu_seqlens.data_ptr<int32_t>();
int64_t max_seq_len = 0;
for (int64_t i = 0; i < batch_size; ++i) {
int64_t span = static_cast<int64_t>(cu_seqlens_ptr[i + 1]) -
static_cast<int64_t>(cu_seqlens_ptr[i]);
TORCH_CHECK(
span >= 0,
"cu_seqlens must be non-decreasing at index ",
i);
if (span > max_seq_len) {
max_seq_len = span;
}
}
TORCH_CHECK(
spec_state_indices.size(1) >= max_seq_len,
"spec_state_indices second dimension must be at least the longest "
"cu_seqlens span; got ",
spec_state_indices.size(1),
" but required ",
max_seq_len);
int64_t spec_stride = spec_state_indices.stride(0);
CHECK_EQ(v_num_heads % num_heads, 0);
TORCH_CHECK(A_log.sizes() == at::IntArrayRef({v_num_heads}));
CHECK_INPUT_SHAPE_DTYPE<true>(
initial_state_source,
{initial_state_source.size(0), v_num_heads, head_dim, v_head_dim},
at::kFloat);
TORCH_CHECK(
initial_state_source.scalar_type() == at::kFloat ||
initial_state_source.scalar_type() == at::kBFloat16 ||
initial_state_source.scalar_type() == at::kHalf,
"fused_sigmoid_gating_delta_rule_update_spec_cpu: initial_state dtype "
"must be float16, bfloat16, or float32");
TORCH_CHECK(
initial_state_source.sizes() ==
at::IntArrayRef(
{initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}),
"Input tensor shape mismatch for initial_state_source");
CHECK_CPU(initial_state_source);
TORCH_CHECK(
initial_state_source.stride(-1) == 1 &&
initial_state_source.stride(-2) == v_head_dim &&
initial_state_source.stride(-3) == head_dim * v_head_dim,
"fused_sigmoid_gating_delta_rule_update_spec_cpu: expect "
"initial_state_source to be contiguous per pool slot.");
TORCH_CHECK(initial_state_source.size(0) >= batch_size,
"initial_state_source capacity too small: size(0)=",
initial_state_source.size(0), ", batch_size=", batch_size);
@@ -2415,39 +2572,56 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(
at::Tensor o = at::empty({total_tokens, v_num_heads, v_head_dim}, q.options());
at::Tensor qk_scale_buf = at::empty({2, total_tokens, num_heads}, at::kFloat);
CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_spec_kernel_impl", [&] {
fused_sigmoid_gating_delta_rule_update_spec_kernel_impl<scalar_t, param_t>(
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
A_log.data_ptr<param_t>(),
a.data_ptr<scalar_t>(),
dt_bias.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
spec_state_indices.data_ptr<int32_t>(),
num_accepted_tokens.data_ptr<int32_t>(),
cu_seqlens.data_ptr<int32_t>(),
initial_state_source.data_ptr<float>(),
o.data_ptr<scalar_t>(),
qk_scale_buf.data_ptr<float>(),
total_tokens,
batch_size,
spec_stride,
num_heads,
head_dim,
v_num_heads,
v_head_dim,
q_strideT,
q_strideH,
k_strideT,
k_strideH,
v_strideT,
v_strideH,
state_slot_stride,
use_qk_l2norm_in_kernel,
softplus_threshold);
});
auto launch = [&](auto state_tag) {
using state_t = decltype(state_tag);
CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
q.scalar_type(), A_log.scalar_type(),
"fused_sigmoid_gating_delta_rule_update_spec_kernel_impl", [&] {
fused_sigmoid_gating_delta_rule_update_spec_kernel_impl<
scalar_t, param_t, state_t>(
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
A_log.data_ptr<param_t>(),
a.data_ptr<scalar_t>(),
dt_bias.data_ptr<scalar_t>(),
b.data_ptr<scalar_t>(),
spec_state_indices.data_ptr<int32_t>(),
num_accepted_tokens.data_ptr<int32_t>(),
cu_seqlens.data_ptr<int32_t>(),
initial_state_source.data_ptr<state_t>(),
o.data_ptr<scalar_t>(),
qk_scale_buf.data_ptr<float>(),
total_tokens,
batch_size,
spec_stride,
num_heads,
head_dim,
v_num_heads,
v_head_dim,
q_strideT,
q_strideH,
k_strideT,
k_strideH,
v_strideT,
v_strideH,
state_slot_stride,
use_qk_l2norm_in_kernel,
softplus_threshold);
});
};
if (initial_state_source.scalar_type() == at::kFloat) {
launch(float());
} else if (initial_state_source.scalar_type() == at::kBFloat16) {
launch(at::BFloat16());
} else if (initial_state_source.scalar_type() == at::kHalf) {
launch(at::Half());
} else {
TORCH_CHECK(
false,
"fused_sigmoid_gating_delta_rule_update_spec_cpu: initial_state "
"dtype must be float16, bfloat16, or float32");
}
return o;
}
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+59 -17
View File
@@ -37,6 +37,7 @@ CHUNK_HEAD_DIMS = [
CHUNK_SIZE = 64
CONV_DIM = 128
CONV_KERNEL = 4
DECODE_BATCH_SIZES = [1, 3, 5]
PREFILL_SEQ_LENS = [
[1],
[1, 2, 3],
@@ -47,7 +48,36 @@ PREFILL_SEQ_LENS = [
[2 * CHUNK_SIZE - 1, 2 * CHUNK_SIZE, 2 * CHUNK_SIZE + 1],
[4 * CHUNK_SIZE + 17],
]
DECODE_BATCH_SIZES = [1, 3, 5]
DECODE_DTYPE_CASES = [
*[
pytest.param(
batch_size,
num_heads,
head_dims,
torch.float32,
id=f"fp32-b{batch_size}-h{num_heads}-d{head_dims}",
)
for batch_size in [1, 3, 5]
for num_heads in NUM_HEADS
for head_dims in HEAD_DIMS
],
pytest.param(3, (2, 4), (32, 32), torch.bfloat16, id="bf16-representative"),
pytest.param(3, (2, 4), (32, 32), torch.float16, id="fp16-representative"),
]
PREFILL_DTYPE_CASES = [
*[
pytest.param(
num_heads,
head_dims,
torch.float32,
id=f"fp32-h{num_heads}-d{head_dims}",
)
for num_heads in NUM_HEADS
for head_dims in CHUNK_HEAD_DIMS
],
pytest.param((2, 4), (64, 64), torch.bfloat16, id="bf16-representative"),
pytest.param((2, 4), (64, 64), torch.float16, id="fp16-representative"),
]
@functools.lru_cache(maxsize=128, typed=False)
@@ -91,6 +121,8 @@ def ref_gated_delta_rule(
initial_state: torch.Tensor,
cu_seqlens: torch.Tensor,
use_qk_l2norm_in_kernel: bool = False,
state_dtype: torch.dtype = torch.float32,
state_write_interval: int = 1,
) -> tuple[torch.Tensor, torch.Tensor]:
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
out = torch.empty_like(value)
@@ -144,6 +176,10 @@ def ref_gated_delta_rule(
delta = (v_t - kv_mem) * beta_t
state = state + delta.unsqueeze(-1) * k_t.unsqueeze(-2)
out_seq[:, :, token_idx] = (state * q_t.unsqueeze(-2)).sum(dim=-1)
if state_dtype != torch.float32 and (
(token_idx + 1) % state_write_interval == 0 or token_idx + 1 == seq_len
):
state = state.to(state_dtype).float()
out[:, begin:end] = out_seq.transpose(1, 2).contiguous().to(initial_dtype)
final_state[seq_idx] = state.squeeze(0)
@@ -201,14 +237,15 @@ def test_fused_gdn_gating_cpu(
# decode path
@pytest.mark.parametrize("batch_size", DECODE_BATCH_SIZES)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_dims", HEAD_DIMS)
@pytest.mark.parametrize(
"batch_size,num_heads,head_dims,state_dtype", DECODE_DTYPE_CASES
)
@torch.inference_mode()
def test_fused_sigmoid_gating_delta_rule_update_cpu(
batch_size: int,
num_heads: tuple[int, int],
head_dims: tuple[int, int],
state_dtype: torch.dtype,
) -> None:
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
num_tokens=batch_size,
@@ -221,7 +258,7 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32)
state_shape = (batch_size, num_v_heads, head_dim, v_head_dim)
state = tensor_cache(
batch_size * num_v_heads * head_dim * v_head_dim, torch.float32
batch_size * num_v_heads * head_dim * v_head_dim, state_dtype
).view(state_shape)
state_ref = state[state_indices].transpose(-1, -2).contiguous()
@@ -236,6 +273,7 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
initial_state=state_ref,
cu_seqlens=cu_seqlens,
use_qk_l2norm_in_kernel=True,
state_dtype=state_dtype,
)
out_ref = out_ref.transpose(0, 1).contiguous()
@@ -256,8 +294,8 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(
state_out[state_indices].transpose(-1, -2),
final_state_ref,
state_out[state_indices].transpose(-1, -2).float(),
final_state_ref.float(),
atol=1e-2,
rtol=1e-2,
)
@@ -265,13 +303,13 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
# prefill path
@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_dims", CHUNK_HEAD_DIMS)
@pytest.mark.parametrize("num_heads,head_dims,state_dtype", PREFILL_DTYPE_CASES)
@torch.inference_mode()
def test_chunk_gated_delta_rule_cpu(
seq_lens: list[int],
num_heads: tuple[int, int],
head_dims: tuple[int, int],
state_dtype: torch.dtype,
) -> None:
total_tokens = sum(seq_lens)
q, k, v, a, b, A_log, dt_bias = gdn_inputs(
@@ -286,7 +324,7 @@ def test_chunk_gated_delta_rule_cpu(
)
initial_state_shape = (len(seq_lens), num_v_heads, head_dim, v_head_dim)
initial_state = tensor_cache(
len(seq_lens) * num_v_heads * head_dim * v_head_dim, torch.float32
len(seq_lens) * num_v_heads * head_dim * v_head_dim, state_dtype
).view(initial_state_shape)
initial_state_ref = initial_state.transpose(-1, -2).contiguous()
@@ -301,6 +339,8 @@ def test_chunk_gated_delta_rule_cpu(
initial_state=initial_state_ref,
cu_seqlens=cu_seqlens,
use_qk_l2norm_in_kernel=True,
state_dtype=state_dtype,
state_write_interval=CHUNK_SIZE,
)
g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
@@ -320,8 +360,8 @@ def test_chunk_gated_delta_rule_cpu(
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(
final_state.transpose(-1, -2),
final_state_ref,
final_state.transpose(-1, -2).float(),
final_state_ref.float(),
atol=1e-2,
rtol=1e-2,
)
@@ -339,14 +379,16 @@ TWO_CALL_SPLITS = [
@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_dims", CHUNK_HEAD_DIMS)
@pytest.mark.parametrize("num_heads", [(2, 4)])
@pytest.mark.parametrize("head_dims", [(64, 64)])
@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16, torch.float16])
@torch.inference_mode()
def test_chunk_gated_delta_rule_cpu_two_call_split(
total_tokens: int,
split: int,
num_heads: tuple[int, int],
head_dims: tuple[int, int],
state_dtype: torch.dtype,
) -> None:
"""A prefill split into two calls (the second seeded with the first's
``final_state`` and a rebased ``cu_seqlens``) must match the single-call
@@ -365,7 +407,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split(
g = g.unsqueeze(0) # [1, T, HV]
beta = beta.unsqueeze(0)
zero_state = torch.zeros(1, num_v_heads, head_dim, v_head_dim, dtype=torch.float32)
zero_state = torch.zeros(1, num_v_heads, head_dim, v_head_dim, dtype=state_dtype)
# Reference: whole sequence in one call, no initial state.
out_full, final_full = ops.chunk_gated_delta_rule_cpu(
@@ -405,7 +447,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split(
value=v[:, split:],
g=g[:, split:],
beta=beta[:, split:],
initial_state=state1.to(torch.float32),
initial_state=state1,
output_final_state=True,
cu_seqlens=torch.tensor([0, tail], dtype=torch.int32),
head_first=False,
@@ -416,7 +458,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split(
out_split = torch.cat([out1, out2], dim=1)
# State must be near-exact; output allows a looser bound for the bf16 round-trip.
torch.testing.assert_close(state2, final_full, atol=1e-3, rtol=1e-3)
torch.testing.assert_close(state2.float(), final_full.float(), atol=1e-3, rtol=1e-3)
torch.testing.assert_close(out_split, out_full, atol=2e-2, rtol=2e-2)
tests/platforms/test_cpu.py
+152
View File
@@ -0,0 +1,152 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
from vllm.config import CacheConfig, VllmConfig
from vllm.platforms.cpu import CpuPlatform
class _FakeModel:
@classmethod
def get_mamba_state_dtype_from_config(
cls, vllm_config: VllmConfig
) -> tuple[torch.dtype, torch.dtype]:
state_dtype = {
"float16": torch.float16,
"bfloat16": torch.bfloat16,
}.get(vllm_config.cache_config.mamba_ssm_cache_dtype, torch.float32)
return torch.float32, state_dtype
class _FakeRegistry:
@staticmethod
def resolve_model_cls(
architecture: str,
model_config: SimpleNamespace,
) -> tuple[type[_FakeModel], str]:
return _FakeModel, architecture
def _cpu_config(
cache_config: CacheConfig,
*,
model_type: str,
resolved_dtype: str,
architecture: str | None = None,
layer_types: tuple[str, ...] = ("linear_attention",),
) -> SimpleNamespace:
model_config = SimpleNamespace(
disable_cascade_attn=False,
architecture=architecture or model_type,
has_inner_state=True,
use_mla=False,
hf_text_config=SimpleNamespace(
model_type=model_type,
layer_types=layer_types,
),
registry=_FakeRegistry,
)
config = VllmConfig(cache_config=cache_config)
config.model_config = model_config
cache_config.mamba_ssm_cache_dtype = resolved_dtype
return config
@pytest.mark.parametrize(
(
"model_type",
"requested_dtype",
"resolved_dtype",
"layer_types",
"expected_dtype",
),
[
pytest.param(
"qwen3_5",
"bfloat16",
"bfloat16",
("linear_attention",),
"bfloat16",
id="gdn-explicit-bf16",
),
pytest.param(
"qwen3_5",
"float16",
"float16",
("linear_attention",),
"float16",
id="gdn-explicit-fp16",
),
pytest.param(
"qwen3_5",
"auto",
"bfloat16",
("linear_attention",),
"bfloat16",
id="gdn-model-bf16",
),
pytest.param(
"qwen3_5",
"auto",
"float16",
("linear_attention",),
"float16",
id="gdn-model-fp16",
),
pytest.param(
"nemotron_h",
"auto",
"bfloat16",
("mamba",),
"float32",
id="unsupported-model-bf16",
),
pytest.param(
"nemotron_h",
"auto",
"float16",
("mamba",),
"float32",
id="unsupported-model-fp16",
),
pytest.param(
"nemotron_h",
"bfloat16",
"bfloat16",
("mamba",),
"float32",
id="unsupported-explicit-bf16",
),
pytest.param(
"nemotron_h",
"float16",
"float16",
("mamba",),
"float32",
id="unsupported-explicit-fp16",
),
],
)
def test_cpu_accelerated_gdn_dtype_policy(
monkeypatch: pytest.MonkeyPatch,
model_type: str,
requested_dtype: str,
resolved_dtype: str,
layer_types: tuple[str, ...],
expected_dtype: str,
) -> None:
monkeypatch.setattr("torch.cpu._is_avx512_bf16_supported", lambda: True)
cache_config = CacheConfig(mamba_ssm_cache_dtype=requested_dtype)
config = _cpu_config(
cache_config,
model_type=model_type,
resolved_dtype=resolved_dtype,
layer_types=layer_types,
)
CpuPlatform.check_and_update_config(config)
assert cache_config.mamba_ssm_cache_dtype == expected_dtype
vllm/platforms/cpu.py
+72 -6
View File
@@ -31,6 +31,61 @@ else:
VllmConfig = None
def _cpu_mamba_backend(vllm_config: VllmConfig) -> str:
"""Return the CPU state backend selected by the model configuration."""
model_config = vllm_config.model_config
if model_config is None:
return "none"
hf_config = getattr(model_config, "hf_text_config", None)
model_type = str(getattr(hf_config, "model_type", "")).lower()
architecture = str(getattr(model_config, "architecture", "")).lower()
layer_types = getattr(hf_config, "layer_types", None)
has_linear_attention = isinstance(layer_types, (list, tuple)) and (
"linear_attention" in layer_types
)
try:
has_inner_state = bool(model_config.has_inner_state)
except (AttributeError, RuntimeError):
has_inner_state = False
if not (has_inner_state or has_linear_attention):
return "none"
fallback_backend = model_type or architecture or "unknown"
if not has_linear_attention:
return fallback_backend
try:
model_cls, _ = model_config.registry.resolve_model_cls(
model_config.architecture,
model_config=model_config,
)
except Exception:
return fallback_backend
try:
state_dtypes = model_cls.get_mamba_state_dtype_from_config(vllm_config)
except Exception:
return fallback_backend
if not isinstance(state_dtypes, tuple) or len(state_dtypes) != 2:
return fallback_backend
if vllm_config.cache_config.mamba_ssm_cache_dtype == "float16":
cache_dtype = torch.float16
elif vllm_config.cache_config.mamba_ssm_cache_dtype == "bfloat16":
cache_dtype = torch.bfloat16
else:
return fallback_backend
if state_dtypes[1] == cache_dtype:
return "gdn"
return fallback_backend
def get_max_threads(pid=0):
if hasattr(os, "sched_getaffinity"):
return len(os.sched_getaffinity(pid))
@@ -181,16 +236,27 @@ class CpuPlatform(Platform):
"otherwise the performance is not optimized."
)
# Accelerated GDN (AMX tiles or AVX-512BF16 VDPBF16PS) requires
# float32 SSM state.
# Accelerated GDN uses AMX tiles or AVX-512BF16 VDPBF16PS.
if (
torch.cpu._is_avx512_bf16_supported()
and cache_config.mamba_ssm_cache_dtype != "float32"
):
cache_config.mamba_ssm_cache_dtype = "float32"
logger.warning(
"Reset SSM cache type to float32 for accelerated GDN mamba attention."
)
mamba_backend = _cpu_mamba_backend(vllm_config)
if (
cache_config.mamba_ssm_cache_dtype in ("float16", "bfloat16")
and mamba_backend == "gdn"
):
logger.info(
"Using %s SSM state storage for the CPU accelerated GDN backend.",
cache_config.mamba_ssm_cache_dtype,
)
else:
cache_config.mamba_ssm_cache_dtype = "float32"
logger.warning(
"Reset SSM cache type to float32 for accelerated GDN mamba "
"backend '%s'.",
mamba_backend,
)
# Lagecy setting
env_key = "VLLM_CPU_KVCACHE_SPACE"