From c6c33f2b1facdcf6ace0d2f241f20507c6cc8157 Mon Sep 17 00:00:00 2001 From: Tianmu Li Date: Mon, 31 Aug 2026 01:32:20 -0700 Subject: [PATCH] [CPU] Support FP16/BF16 persisted GDN state on AMX (#52191) Signed-off-by: Li, Tianmu Co-authored-by: OpenAI Codex --- csrc/cpu/sgl-kernels/fla.cpp | 410 ++++++++++++++------ tests/kernels/mamba/cpu/test_cpu_gdn_ops.py | 76 +++- tests/platforms/test_cpu.py | 152 ++++++++ vllm/platforms/cpu.py | 78 +++- 4 files changed, 575 insertions(+), 141 deletions(-) create mode 100644 tests/platforms/test_cpu.py diff --git a/csrc/cpu/sgl-kernels/fla.cpp b/csrc/cpu/sgl-kernels/fla.cpp index fec6648cec..dd6f3b29ab 100644 --- a/csrc/cpu/sgl-kernels/fla.cpp +++ b/csrc/cpu/sgl-kernels/fla.cpp @@ -74,6 +74,20 @@ void pack_vnni2(scalar_t* __restrict__ dst, float* __restrict__ src, const float #endif } +template +inline void store_state_vec2( + state_t* __restrict__ dst, + const at::vec::Vectorized& src0, + const at::vec::Vectorized& src1) { + if constexpr (std::is_same_v) { + src0.store(dst); + src1.store(dst + at::vec::Vectorized::size()); + } else { + auto state_vec = convert_from_float_ext(src0, src1); + state_vec.store(dst); + } +} + template inline void fill_stub(scalar_t* __restrict__ out, float val) { using Vec = at::vec::Vectorized; @@ -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 +template 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::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) { + s_ptr = state_slot; + } else { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + 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(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) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + 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_slot + d0, state_vec0, state_vec1); + } + for (; d0 < D * D; ++d0) { + state_slot[d0] = static_cast(s_ptr[d0]); + } + } } // move to the next index @@ -1576,7 +1627,7 @@ inline at::vec::Vectorized softplus(const at::vec::Vectorized& x, return Vec::blendv(Vec::blendv(log1pex, expx, mask_lo), x, mask_hi); } -template +template 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_ptr + state_offset + di * v_head_dim + dvi, + state_vec0, + state_vec1); } bVec o_vec = at::vec::convert_from_float(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(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(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_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 +template 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; using fVec = at::vec::Vectorized; 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( + dst + di * v_head_dim + dvi, sv0, sv1); } bVec o_vec = at::vec::convert_from_float(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(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(src[di * v_head_dim + dvi]) * + g_val_exp + + k_val * dt_val; + dst[di * v_head_dim + dvi] = static_cast(ns); o_val += ns * q_val * scale; } o_ptr[o_offset + dvi] = (scalar_t)o_val; @@ -2112,10 +2176,10 @@ std::tuple 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( \ +#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \ + chunk_gated_delta_rule_fwd_inter_kernel_impl( \ o.data_ptr(), \ - initial_state.data_ptr(), \ + initial_state.data_ptr(), \ initial_state_indices.data_ptr(), \ q.data_ptr(), \ k.data_ptr(), \ @@ -2134,7 +2198,7 @@ std::tuple chunk_gated_delta_rule_fwd_intra( k.stride(2), \ initial_state.stride(0)); -template +template std::tuple chunk_gated_delta_rule_fwd_inter( const at::Tensor& q, const at::Tensor& k, @@ -2172,7 +2236,7 @@ std::tuple 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 chunk_gated_delta_rule_cpu( @@ -2211,7 +2275,12 @@ std::tuple chunk_gated_delta_rule_cpu( CHECK_INPUT_SHAPE_DTYPE(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 chunk_gated_delta_rule_cpu( chunk_gated_delta_rule_fwd_intra(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( - 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( + 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 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(b, {batch_size, v_num_heads}, q.scalar_type()); CHECK_INPUT_SHAPE_DTYPE(initial_state_indices, {batch_size}, at::kInt); CHECK_INPUT_SHAPE_DTYPE(cu_seqlens, {batch_size + 1}, at::kInt); - CHECK_INPUT_SHAPE_DTYPE( - 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( - q.data_ptr(), - k.data_ptr(), - v.data_ptr(), - A_log.data_ptr(), - a.data_ptr(), - dt_bias.data_ptr(), - b.data_ptr(), - initial_state_indices.data_ptr(), - initial_state_source.data_ptr(), - core_attn_out.data_ptr(), - qk_scale_buf.data_ptr(), - 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( + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + A_log.data_ptr(), + a.data_ptr(), + dt_bias.data_ptr(), + b.data_ptr(), + initial_state_indices.data_ptr(), + initial_state_source.data_ptr(), + core_attn_out.data_ptr(), + qk_scale_buf.data_ptr(), + 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(k, {total_tokens, num_heads, head_dim}, q.scalar_type()); CHECK_INPUT_SHAPE_DTYPE(v, {total_tokens, v_num_heads, v_head_dim}, q.scalar_type()); CHECK_INPUT_SHAPE_DTYPE(a, {total_tokens, v_num_heads}, q.scalar_type()); CHECK_INPUT_SHAPE_DTYPE(b, {total_tokens, v_num_heads}, q.scalar_type()); CHECK_INPUT_SHAPE_DTYPE(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(num_accepted_tokens, {batch_size}, at::kInt); CHECK_INPUT_SHAPE_DTYPE(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(); + int64_t max_seq_len = 0; + for (int64_t i = 0; i < batch_size; ++i) { + int64_t span = static_cast(cu_seqlens_ptr[i + 1]) - + static_cast(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( - 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( - q.data_ptr(), - k.data_ptr(), - v.data_ptr(), - A_log.data_ptr(), - a.data_ptr(), - dt_bias.data_ptr(), - b.data_ptr(), - spec_state_indices.data_ptr(), - num_accepted_tokens.data_ptr(), - cu_seqlens.data_ptr(), - initial_state_source.data_ptr(), - o.data_ptr(), - qk_scale_buf.data_ptr(), - 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(), + k.data_ptr(), + v.data_ptr(), + A_log.data_ptr(), + a.data_ptr(), + dt_bias.data_ptr(), + b.data_ptr(), + spec_state_indices.data_ptr(), + num_accepted_tokens.data_ptr(), + cu_seqlens.data_ptr(), + initial_state_source.data_ptr(), + o.data_ptr(), + qk_scale_buf.data_ptr(), + 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; } diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index 385870e49c..c239f831c6 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -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) diff --git a/tests/platforms/test_cpu.py b/tests/platforms/test_cpu.py new file mode 100644 index 0000000000..fd270315eb --- /dev/null +++ b/tests/platforms/test_cpu.py @@ -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 diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 724f840129..b0b5228829 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -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"