@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c6c33f2b1f

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

Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Tianmu Li committed Aug 31, 2026 at 01:32 UTC c6c33f2b1facdcf6ace0d2f241f20507c6cc8157
4 files changed +575 -141
csrc/cpu/sgl-kernels/fla.cpp
+292 -118
@@ -74,6 +74,20 @@ void pack_vnni2(scalar_t* __restrict__ dst, float* __restrict__ src, const float
74 #endif
75 }
76
77 +template <typename state_t>
78 +inline void store_state_vec2(
79 + state_t* __restrict__ dst,
80 + const at::vec::Vectorized<float>& src0,
81 + const at::vec::Vectorized<float>& src1) {
82 + if constexpr (std::is_same_v<state_t, float>) {
83 + src0.store(dst);
84 + src1.store(dst + at::vec::Vectorized<float>::size());
85 + } else {
86 + auto state_vec = convert_from_float_ext<state_t>(src0, src1);
87 + state_vec.store(dst);
88 + }
89 +}
90 +
91 template <typename scalar_t, int SIZE>
92 inline void fill_stub(scalar_t* __restrict__ out, float val) {
93 using Vec = at::vec::Vectorized<scalar_t>;
@@ -1255,10 +1269,10 @@ void chunk_gated_delta_rule_fwd_intra_kernel_impl(
1269 // d : [B, NT, Hv, C, C]
1270 // cu_seqlens : [num_seqs + 1]
1271 // chunk_offsets : [num_seqs + 1]
1258 -template <typename scalar_t, int D, int CHUNK_SIZE>
1272 +template <typename scalar_t, typename state_t, int D, int CHUNK_SIZE>
1273 void chunk_gated_delta_rule_fwd_inter_kernel_impl(
1274 scalar_t* __restrict__ out,
1261 - float* __restrict__ state,
1275 + state_t* __restrict__ state,
1276 const int32_t* __restrict__ indices,
1277 const scalar_t* __restrict__ q,
1278 const scalar_t* __restrict__ k,
@@ -1297,6 +1311,7 @@ void chunk_gated_delta_rule_fwd_inter_kernel_impl(
1311 // thread local temp buffer
1312 DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D);
1313 DECL_ZERO_BUF(scalar_t, tmp2, D * D);
1314 + DECL_ZERO_BUF(float, state_tmp, D * D);
1315 DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D);
1316 DECL_ZERO_BUF(scalar_t, tmp4, CHUNK_SIZE * D);
1317 DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE);
@@ -1370,7 +1385,27 @@ void chunk_gated_delta_rule_fwd_inter_kernel_impl(
1385 apply_mask_kernel<scalar_t, CHUNK_SIZE, false>::apply(attn2, attn, nullptr, d_ptr, mb_size);
1386
1387 // step 2.a: v' = w @ state (fuse state *= exp(g_last) with packing)
1373 - float* __restrict__ s_ptr = state + indices[bs] * state_strideS + hv * (D * D);
1388 + state_t* __restrict__ state_slot =
1389 + state + indices[bs] * state_strideS + hv * (D * D);
1390 + float* s_ptr;
1391 + if constexpr (std::is_same_v<state_t, float>) {
1392 + s_ptr = state_slot;
1393 + } else {
1394 + using bVec = at::vec::Vectorized<state_t>;
1395 + using fVec = at::vec::Vectorized<float>;
1396 + constexpr int kVecSize = bVec::size();
1397 + constexpr int fVecSize = fVec::size();
1398 + int64_t d0 = 0;
1399 + for (; d0 + kVecSize <= D * D; d0 += kVecSize) {
1400 + auto [state_vec0, state_vec1] = load_float_vec2(state_slot + d0);
1401 + state_vec0.store(state_tmp + d0);
1402 + state_vec1.store(state_tmp + d0 + fVecSize);
1403 + }
1404 + for (; d0 < D * D; ++d0) {
1405 + state_tmp[d0] = static_cast<float>(state_slot[d0]);
1406 + }
1407 + s_ptr = state_tmp;
1408 + }
1409 const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * (CHUNK_SIZE);
1410 float g_last = g_ptr[mb_size - 1];
1411 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(
1582 s_ptr,
1583 D);
1584 }
1585 +
1586 + if constexpr (!std::is_same_v<state_t, float>) {
1587 + using bVec = at::vec::Vectorized<state_t>;
1588 + using fVec = at::vec::Vectorized<float>;
1589 + constexpr int kVecSize = bVec::size();
1590 + constexpr int fVecSize = fVec::size();
1591 + int64_t d0 = 0;
1592 + for (; d0 + kVecSize <= D * D; d0 += kVecSize) {
1593 + fVec state_vec0 = fVec::loadu(s_ptr + d0);
1594 + fVec state_vec1 = fVec::loadu(s_ptr + d0 + fVecSize);
1595 + store_state_vec2<state_t>(state_slot + d0, state_vec0, state_vec1);
1596 + }
1597 + for (; d0 < D * D; ++d0) {
1598 + state_slot[d0] = static_cast<state_t>(s_ptr[d0]);
1599 + }
1600 + }
1601 }
1602
1603 // move to the next index
@@ -1576,7 +1627,7 @@ inline at::vec::Vectorized<float> softplus(const at::vec::Vectorized<float>& x,
1627 return Vec::blendv(Vec::blendv(log1pex, expx, mask_lo), x, mask_hi);
1628 }
1629
1579 -template <typename scalar_t, typename param_t>
1630 +template <typename scalar_t, typename param_t, typename state_t>
1631 void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1632 const scalar_t* __restrict__ q_ptr,
1633 const scalar_t* __restrict__ k_ptr,
@@ -1586,7 +1637,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1637 const scalar_t* __restrict__ dt_bias_ptr,
1638 const scalar_t* __restrict__ b_ptr,
1639 const int32_t* __restrict__ indices_ptr,
1589 - float* __restrict__ state_ptr,
1640 + state_t* __restrict__ state_ptr,
1641 scalar_t* __restrict__ o_ptr,
1642 float* __restrict__ qk_scale_buf,
1643 int64_t seq_len,
@@ -1700,8 +1751,10 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1751 state_vec1 = state_vec1 * g_val_exp_vec + k_vec * dt_vec1;
1752 o_vec0 = o_vec0 + state_vec0 * q_vec * scale_vec;
1753 o_vec1 = o_vec1 + state_vec1 * q_vec * scale_vec;
1703 - state_vec0.store(state_ptr + state_offset + di * v_head_dim + dvi);
1704 - state_vec1.store(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
1754 + store_state_vec2<state_t>(
1755 + state_ptr + state_offset + di * v_head_dim + dvi,
1756 + state_vec0,
1757 + state_vec1);
1758 }
1759 bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1);
1760 o_vec.store(o_ptr + o_offset + dvi);
@@ -1710,8 +1763,10 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1763 float kv_mem_val = 0;
1764 for (int di = 0; di < head_dim; ++di) {
1765 float k_val = k_ptr[k_offset + di] * k_scale;
1713 - state_ptr[state_offset + di * v_head_dim + dvi] *= g_val_exp;
1714 - kv_mem_val += state_ptr[state_offset + di * v_head_dim + dvi] * k_val;
1766 + float state_val =
1767 + static_cast<float>(state_ptr[state_offset + di * v_head_dim + dvi]) *
1768 + g_val_exp;
1769 + kv_mem_val += state_val * k_val;
1770 }
1771 float v_val = v_ptr[v_offset + dvi];
1772 float dt_val = (v_val - kv_mem_val) * beta_val;
@@ -1719,8 +1774,13 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1774 for (int di = 0; di < head_dim; ++di) {
1775 float q_val = q_ptr[q_offset + di] * q_scale;
1776 float k_val = k_ptr[k_offset + di] * k_scale;
1722 - state_ptr[state_offset + di * v_head_dim + dvi] += k_val * dt_val;
1723 - o_val += state_ptr[state_offset + di * v_head_dim + dvi] * q_val * scale;
1777 + float state_val =
1778 + static_cast<float>(state_ptr[state_offset + di * v_head_dim + dvi]) *
1779 + g_val_exp;
1780 + state_val += k_val * dt_val;
1781 + state_ptr[state_offset + di * v_head_dim + dvi] =
1782 + static_cast<state_t>(state_val);
1783 + o_val += state_val * q_val * scale;
1784 }
1785 o_ptr[o_offset + dvi] = o_val;
1786 }
@@ -1736,7 +1796,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1796 // token ``t`` into cache slot ``t`` (multi-slot rollback, matching the GPU
1797 // kernel). Parallelized over (sequence, v_head); the per-sequence token loop is
1798 // sequential as required by the recurrence.
1739 -template <typename scalar_t, typename param_t>
1799 +template <typename scalar_t, typename param_t, typename state_t>
1800 void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
1801 const scalar_t* __restrict__ q_ptr, // [T, HK, EK]
1802 const scalar_t* __restrict__ k_ptr, // [T, HK, EK]
@@ -1748,7 +1808,7 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
1808 const int32_t* __restrict__ spec_indices_ptr, // [N, S]
1809 const int32_t* __restrict__ num_accepted_ptr, // [N]
1810 const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1]
1751 - float* __restrict__ state_ptr,
1811 + state_t* __restrict__ state_ptr,
1812 scalar_t* __restrict__ o_ptr, // [T, HV, EV]
1813 float* __restrict__ qk_scale_buf, // [2, T, HK]
1814 int64_t total_tokens,
@@ -1770,7 +1830,6 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
1830 using bVec = at::vec::Vectorized<scalar_t>;
1831 using fVec = at::vec::Vectorized<float>;
1832 constexpr int64_t VecSize = bVec::size();
1773 - constexpr int64_t fVecSize = fVec::size();
1833 int64_t group_size = v_num_heads / num_heads;
1834 double scale = 1 / std::sqrt((double)head_dim);
1835 fVec scale_vec = fVec((float)scale);
@@ -1815,8 +1874,10 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
1874 for (int64_t t = 0; t < q_len; ++t) {
1875 int64_t cur_slot = (int64_t)spec_indices_ptr[bi * spec_stride + t];
1876 int64_t token = q_start + t;
1818 - const float* src = state_ptr + prev_slot * state_slot_stride + ni * head_dim * v_head_dim;
1819 - float* dst = state_ptr + cur_slot * state_slot_stride + ni * head_dim * v_head_dim;
1877 + const state_t* src =
1878 + state_ptr + prev_slot * state_slot_stride + ni * head_dim * v_head_dim;
1879 + state_t* dst =
1880 + state_ptr + cur_slot * state_slot_stride + ni * head_dim * v_head_dim;
1881 float g_val = -std::exp((float)A_log_ptr[ni]) *
1882 softplus((float)a_ptr[token * v_num_heads + ni] + (float)dt_bias_ptr[ni], softplus_threshold);
1883 float g_val_exp = std::exp(g_val);
@@ -1836,8 +1897,8 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
1897 fVec kv_mem_vec1 = fVec(0.f);
1898 for (int di = 0; di < head_dim; ++di) {
1899 fVec k_val_vec = fVec((float)k_ptr[k_offset + di] * k_scale);
1839 - fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi);
1840 - fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize);
1900 + auto [sv0, sv1] =
1901 + load_float_vec2(src + di * v_head_dim + dvi);
1902 kv_mem_vec0 = kv_mem_vec0 + sv0 * g_val_exp_vec * k_val_vec;
1903 kv_mem_vec1 = kv_mem_vec1 + sv1 * g_val_exp_vec * k_val_vec;
1904 }
@@ -1851,14 +1912,14 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
1912 for (int di = 0; di < head_dim; ++di) {
1913 fVec q_vec = fVec((float)q_ptr[q_offset + di] * q_scale);
1914 fVec k_vec = fVec((float)k_ptr[k_offset + di] * k_scale);
1854 - fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi);
1855 - fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize);
1915 + auto [sv0, sv1] =
1916 + load_float_vec2(src + di * v_head_dim + dvi);
1917 sv0 = sv0 * g_val_exp_vec + k_vec * dt_vec0;
1918 sv1 = sv1 * g_val_exp_vec + k_vec * dt_vec1;
1919 o_vec0 = o_vec0 + sv0 * q_vec * scale_vec;
1920 o_vec1 = o_vec1 + sv1 * q_vec * scale_vec;
1860 - sv0.store(dst + di * v_head_dim + dvi);
1861 - sv1.store(dst + di * v_head_dim + dvi + fVecSize);
1921 + store_state_vec2<state_t>(
1922 + dst + di * v_head_dim + dvi, sv0, sv1);
1923 }
1924 bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1);
1925 o_vec.store(o_ptr + o_offset + dvi);
@@ -1867,7 +1928,8 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
1928 float kv_mem_val = 0.f;
1929 for (int di = 0; di < head_dim; ++di) {
1930 float k_val = (float)k_ptr[k_offset + di] * k_scale;
1870 - kv_mem_val += src[di * v_head_dim + dvi] * g_val_exp * k_val;
1931 + kv_mem_val += static_cast<float>(src[di * v_head_dim + dvi]) *
1932 + g_val_exp * k_val;
1933 }
1934 float v_val = (float)v_ptr[v_offset + dvi];
1935 float dt_val = (v_val - kv_mem_val) * beta_val;
@@ -1875,8 +1937,10 @@ void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
1937 for (int di = 0; di < head_dim; ++di) {
1938 float q_val = (float)q_ptr[q_offset + di] * q_scale;
1939 float k_val = (float)k_ptr[k_offset + di] * k_scale;
1878 - float ns = src[di * v_head_dim + dvi] * g_val_exp + k_val * dt_val;
1879 - dst[di * v_head_dim + dvi] = ns;
1940 + float ns = static_cast<float>(src[di * v_head_dim + dvi]) *
1941 + g_val_exp +
1942 + k_val * dt_val;
1943 + dst[di * v_head_dim + dvi] = static_cast<state_t>(ns);
1944 o_val += ns * q_val * scale;
1945 }
1946 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(
2176 return std::make_tuple(w, u, decay_mask);
2177 }
2178
2115 -#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \
2116 - chunk_gated_delta_rule_fwd_inter_kernel_impl<scalar_t, HD, CHUNK_SIZE>( \
2179 +#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \
2180 + chunk_gated_delta_rule_fwd_inter_kernel_impl<scalar_t, state_t, HD, CHUNK_SIZE>( \
2181 o.data_ptr<scalar_t>(), \
2118 - initial_state.data_ptr<float>(), \
2182 + initial_state.data_ptr<state_t>(), \
2183 initial_state_indices.data_ptr<int32_t>(), \
2184 q.data_ptr<scalar_t>(), \
2185 k.data_ptr<scalar_t>(), \
@@ -2134,7 +2198,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_intra(
2198 k.stride(2), \
2199 initial_state.stride(0));
2200
2137 -template <int CHUNK_SIZE>
2201 +template <typename state_t, int CHUNK_SIZE>
2202 std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_inter(
2203 const at::Tensor& q,
2204 const at::Tensor& k,
@@ -2172,7 +2236,7 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_inter(
2236 // value: [B, T, Hv, Dv]
2237 // g: [B, T, Hv] FP32
2238 // beta: [B, T, Hv]
2175 -// initial_state: [num_seqs, Hv, Dv, D] FP32
2239 +// initial_state: [num_seqs, Hv, Dv, D] FP16, BF16, or FP32 persisted state
2240 // cu_seqlens: [num_seqs + 1] INT32
2241 //
2242 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(
2275 CHECK_INPUT_SHAPE_DTYPE<false>(cu_seqlens, {num_seqs + 1}, at::kInt);
2276 TORCH_CHECK(initial_state.sizes() == at::IntArrayRef({initial_state.size(0), Hv, Dv, D}),
2277 "chunk_gated_delta_rule_cpu: initial_state shape mismatch, got ", initial_state.sizes());
2214 - TORCH_CHECK(initial_state.scalar_type() == at::kFloat, "chunk_gated_delta_rule_cpu: initial_state dtype mismatch");
2278 + TORCH_CHECK(
2279 + initial_state.scalar_type() == at::kFloat ||
2280 + initial_state.scalar_type() == at::kBFloat16 ||
2281 + initial_state.scalar_type() == at::kHalf,
2282 + "chunk_gated_delta_rule_cpu: initial_state dtype must be float16, "
2283 + "bfloat16, or float32");
2284 CHECK_CPU(initial_state);
2285 // initial_state may be a pooled/paged buffer with padding between slots
2286 // (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(
2306 chunk_gated_delta_rule_fwd_intra<CHUNK_SIZE>(key_, value, g_, beta, cu_seqlens, chunk_indices);
2307
2308 // fused `chunk_gated_delta_rule_fwd_h` + `chunk_fwd_o`
2240 - auto [output, final_state] = chunk_gated_delta_rule_fwd_inter<CHUNK_SIZE>(
2241 - query_,
2242 - key_,
2243 - w,
2244 - u,
2245 - g_,
2246 - decay_mask,
2247 - initial_state,
2248 - output_final_state,
2249 - cu_seqlens,
2250 - chunk_offsets,
2251 - initial_state_indices);
2309 + auto launch_inter = [&](auto state_tag) {
2310 + using state_t = decltype(state_tag);
2311 + return chunk_gated_delta_rule_fwd_inter<state_t, CHUNK_SIZE>(
2312 + query_,
2313 + key_,
2314 + w,
2315 + u,
2316 + g_,
2317 + decay_mask,
2318 + initial_state,
2319 + output_final_state,
2320 + cu_seqlens,
2321 + chunk_offsets,
2322 + initial_state_indices);
2323 + };
2324 + auto output_and_state =
2325 + initial_state.scalar_type() == at::kFloat
2326 + ? launch_inter(float{})
2327 + : initial_state.scalar_type() == at::kBFloat16
2328 + ? launch_inter(at::BFloat16{})
2329 + : launch_inter(at::Half{});
2330 + auto [output, final_state] = output_and_state;
2331
2332 return std::make_tuple(output, final_state);
2333 }
@@ -2261,6 +2340,7 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
2340 // a: [batch_size, v_num_heads]
2341 // b: [batch_size, v_num_heads]
2342 // initial_state_source:[num_tokens, v_num_heads, head_dim, v_head_dim]
2343 +// FP16, BF16, or FP32 persisted state; recurrent arithmetic remains FP32.
2344 // initial_state_indices: [batch_size]
2345 // cu_seqlens: [batch_size + 1]
2346 at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
@@ -2293,8 +2373,24 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
2373 CHECK_INPUT_SHAPE_DTYPE<true>(b, {batch_size, v_num_heads}, q.scalar_type());
2374 CHECK_INPUT_SHAPE_DTYPE<true>(initial_state_indices, {batch_size}, at::kInt);
2375 CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt);
2296 - CHECK_INPUT_SHAPE_DTYPE<true>(
2297 - initial_state_source, {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}, at::kFloat);
2376 + TORCH_CHECK(
2377 + initial_state_source.scalar_type() == at::kFloat ||
2378 + initial_state_source.scalar_type() == at::kBFloat16 ||
2379 + initial_state_source.scalar_type() == at::kHalf,
2380 + "fused_sigmoid_gating_delta_rule_update_cpu: initial_state dtype must "
2381 + "be float16, bfloat16, or float32");
2382 + TORCH_CHECK(
2383 + initial_state_source.sizes() ==
2384 + at::IntArrayRef(
2385 + {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}),
2386 + "Input tensor shape mismatch for initial_state_source");
2387 + CHECK_CPU(initial_state_source);
2388 + TORCH_CHECK(
2389 + initial_state_source.stride(-1) == 1 &&
2390 + initial_state_source.stride(-2) == v_head_dim &&
2391 + initial_state_source.stride(-3) == head_dim * v_head_dim,
2392 + "fused_sigmoid_gating_delta_rule_update_cpu: expect initial_state_source "
2393 + "to be contiguous per pool slot.");
2394 CHECK(initial_state_source.size(0) >= batch_size);
2395 CHECK_EQ(v_num_heads % num_heads, 0);
2396 TORCH_CHECK(
@@ -2318,46 +2414,63 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
2414 at::Tensor core_attn_out = at::empty({batch_size, seq_len, v_num_heads, v_head_dim}, q.options());
2415 at::Tensor qk_scale_buf = at::empty({2 * batch_size, seq_len, num_heads}, at::kFloat);
2416
2321 - CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
2322 - q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] {
2323 - fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t, param_t>(
2324 - q.data_ptr<scalar_t>(),
2325 - k.data_ptr<scalar_t>(),
2326 - v.data_ptr<scalar_t>(),
2327 - A_log.data_ptr<param_t>(),
2328 - a.data_ptr<scalar_t>(),
2329 - dt_bias.data_ptr<scalar_t>(),
2330 - b.data_ptr<scalar_t>(),
2331 - initial_state_indices.data_ptr<int32_t>(),
2332 - initial_state_source.data_ptr<float>(),
2333 - core_attn_out.data_ptr<scalar_t>(),
2334 - qk_scale_buf.data_ptr<float>(),
2335 - seq_len,
2336 - batch_size,
2337 - num_heads,
2338 - head_dim,
2339 - v_num_heads,
2340 - v_head_dim,
2341 - q_strideB,
2342 - q_strideS,
2343 - q_strideH,
2344 - k_strideB,
2345 - k_strideS,
2346 - k_strideH,
2347 - v_strideB,
2348 - v_strideS,
2349 - v_strideH,
2350 - state_slot_stride,
2351 - use_qk_l2norm_in_kernel,
2352 - softplus_threshold);
2353 - });
2417 + auto launch = [&](auto state_tag) {
2418 + using state_t = decltype(state_tag);
2419 + CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
2420 + q.scalar_type(), A_log.scalar_type(),
2421 + "fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] {
2422 + fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t, param_t, state_t>(
2423 + q.data_ptr<scalar_t>(),
2424 + k.data_ptr<scalar_t>(),
2425 + v.data_ptr<scalar_t>(),
2426 + A_log.data_ptr<param_t>(),
2427 + a.data_ptr<scalar_t>(),
2428 + dt_bias.data_ptr<scalar_t>(),
2429 + b.data_ptr<scalar_t>(),
2430 + initial_state_indices.data_ptr<int32_t>(),
2431 + initial_state_source.data_ptr<state_t>(),
2432 + core_attn_out.data_ptr<scalar_t>(),
2433 + qk_scale_buf.data_ptr<float>(),
2434 + seq_len,
2435 + batch_size,
2436 + num_heads,
2437 + head_dim,
2438 + v_num_heads,
2439 + v_head_dim,
2440 + q_strideB,
2441 + q_strideS,
2442 + q_strideH,
2443 + k_strideB,
2444 + k_strideS,
2445 + k_strideH,
2446 + v_strideB,
2447 + v_strideS,
2448 + v_strideH,
2449 + state_slot_stride,
2450 + use_qk_l2norm_in_kernel,
2451 + softplus_threshold);
2452 + });
2453 + };
2454 + if (initial_state_source.scalar_type() == at::kFloat) {
2455 + launch(float());
2456 + } else if (initial_state_source.scalar_type() == at::kBFloat16) {
2457 + launch(at::BFloat16());
2458 + } else if (initial_state_source.scalar_type() == at::kHalf) {
2459 + launch(at::Half());
2460 + } else {
2461 + TORCH_CHECK(
2462 + false,
2463 + "fused_sigmoid_gating_delta_rule_update_cpu: initial_state dtype must "
2464 + "be float16, bfloat16, or float32");
2465 + }
2466 return core_attn_out;
2467 }
2468
2469 // Speculative-decode update (multi-token, multi-slot rollback).
2470 // q: [T, HK, EK] k: [T, HK, EK] v: [T, HV, EV]
2471 // a: [T, HV] b: [T, HV]
2360 -// initial_state_source: [N_slots, HV, EK, EV] FP32 (updated in place)
2472 +// initial_state_source: [N_slots, HV, EK, EV] FP16, BF16, or FP32 persisted state
2473 +// (updated in place; recurrent arithmetic remains FP32)
2474 // spec_state_indices: [batch, S] INT32 (S = num_spec + 1)
2475 // num_accepted_tokens: [batch] INT32
2476 // cu_seqlens: [batch + 1] INT32
@@ -2385,21 +2498,65 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(
2498 int64_t head_dim = q.size(2);
2499 int64_t v_num_heads = v.size(1);
2500 int64_t v_head_dim = v.size(2);
2388 - int64_t batch_size = cu_seqlens.size(0) - 1;
2389 - int64_t spec_stride = spec_state_indices.stride(0);
2501 CHECK_INPUT_SHAPE_DTYPE<true>(k, {total_tokens, num_heads, head_dim}, q.scalar_type());
2502 CHECK_INPUT_SHAPE_DTYPE<true>(v, {total_tokens, v_num_heads, v_head_dim}, q.scalar_type());
2503 CHECK_INPUT_SHAPE_DTYPE<true>(a, {total_tokens, v_num_heads}, q.scalar_type());
2504 CHECK_INPUT_SHAPE_DTYPE<true>(b, {total_tokens, v_num_heads}, q.scalar_type());
2505 CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type());
2506 + CHECK_DIM(1, cu_seqlens);
2507 + int64_t batch_size = cu_seqlens.size(0) - 1;
2508 CHECK_INPUT_SHAPE_DTYPE<true>(num_accepted_tokens, {batch_size}, at::kInt);
2509 CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt);
2510 + CHECK_DIM(2, spec_state_indices);
2511 + CHECK_CPU(spec_state_indices);
2512 + TORCH_CHECK(
2513 + spec_state_indices.scalar_type() == at::kInt,
2514 + "spec_state_indices must be an int32 CPU tensor");
2515 + TORCH_CHECK(
2516 + spec_state_indices.size(0) == batch_size,
2517 + "spec_state_indices first dimension must equal batch_size");
2518 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(spec_state_indices);
2519 + const int32_t* cu_seqlens_ptr = cu_seqlens.data_ptr<int32_t>();
2520 + int64_t max_seq_len = 0;
2521 + for (int64_t i = 0; i < batch_size; ++i) {
2522 + int64_t span = static_cast<int64_t>(cu_seqlens_ptr[i + 1]) -
2523 + static_cast<int64_t>(cu_seqlens_ptr[i]);
2524 + TORCH_CHECK(
2525 + span >= 0,
2526 + "cu_seqlens must be non-decreasing at index ",
2527 + i);
2528 + if (span > max_seq_len) {
2529 + max_seq_len = span;
2530 + }
2531 + }
2532 + TORCH_CHECK(
2533 + spec_state_indices.size(1) >= max_seq_len,
2534 + "spec_state_indices second dimension must be at least the longest "
2535 + "cu_seqlens span; got ",
2536 + spec_state_indices.size(1),
2537 + " but required ",
2538 + max_seq_len);
2539 + int64_t spec_stride = spec_state_indices.stride(0);
2540 CHECK_EQ(v_num_heads % num_heads, 0);
2541 TORCH_CHECK(A_log.sizes() == at::IntArrayRef({v_num_heads}));
2399 - CHECK_INPUT_SHAPE_DTYPE<true>(
2400 - initial_state_source,
2401 - {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim},
2402 - at::kFloat);
2542 + TORCH_CHECK(
2543 + initial_state_source.scalar_type() == at::kFloat ||
2544 + initial_state_source.scalar_type() == at::kBFloat16 ||
2545 + initial_state_source.scalar_type() == at::kHalf,
2546 + "fused_sigmoid_gating_delta_rule_update_spec_cpu: initial_state dtype "
2547 + "must be float16, bfloat16, or float32");
2548 + TORCH_CHECK(
2549 + initial_state_source.sizes() ==
2550 + at::IntArrayRef(
2551 + {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}),
2552 + "Input tensor shape mismatch for initial_state_source");
2553 + CHECK_CPU(initial_state_source);
2554 + TORCH_CHECK(
2555 + initial_state_source.stride(-1) == 1 &&
2556 + initial_state_source.stride(-2) == v_head_dim &&
2557 + initial_state_source.stride(-3) == head_dim * v_head_dim,
2558 + "fused_sigmoid_gating_delta_rule_update_spec_cpu: expect "
2559 + "initial_state_source to be contiguous per pool slot.");
2560 TORCH_CHECK(initial_state_source.size(0) >= batch_size,
2561 "initial_state_source capacity too small: size(0)=",
2562 initial_state_source.size(0), ", batch_size=", batch_size);
@@ -2415,39 +2572,56 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(
2572 at::Tensor o = at::empty({total_tokens, v_num_heads, v_head_dim}, q.options());
2573 at::Tensor qk_scale_buf = at::empty({2, total_tokens, num_heads}, at::kFloat);
2574
2418 - CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
2419 - q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_spec_kernel_impl", [&] {
2420 - fused_sigmoid_gating_delta_rule_update_spec_kernel_impl<scalar_t, param_t>(
2421 - q.data_ptr<scalar_t>(),
2422 - k.data_ptr<scalar_t>(),
2423 - v.data_ptr<scalar_t>(),
2424 - A_log.data_ptr<param_t>(),
2425 - a.data_ptr<scalar_t>(),
2426 - dt_bias.data_ptr<scalar_t>(),
2427 - b.data_ptr<scalar_t>(),
2428 - spec_state_indices.data_ptr<int32_t>(),
2429 - num_accepted_tokens.data_ptr<int32_t>(),
2430 - cu_seqlens.data_ptr<int32_t>(),
2431 - initial_state_source.data_ptr<float>(),
2432 - o.data_ptr<scalar_t>(),
2433 - qk_scale_buf.data_ptr<float>(),
2434 - total_tokens,
2435 - batch_size,
2436 - spec_stride,
2437 - num_heads,
2438 - head_dim,
2439 - v_num_heads,
2440 - v_head_dim,
2441 - q_strideT,
2442 - q_strideH,
2443 - k_strideT,
2444 - k_strideH,
2445 - v_strideT,
2446 - v_strideH,
2447 - state_slot_stride,
2448 - use_qk_l2norm_in_kernel,
2449 - softplus_threshold);
2450 - });
2575 + auto launch = [&](auto state_tag) {
2576 + using state_t = decltype(state_tag);
2577 + CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
2578 + q.scalar_type(), A_log.scalar_type(),
2579 + "fused_sigmoid_gating_delta_rule_update_spec_kernel_impl", [&] {
2580 + fused_sigmoid_gating_delta_rule_update_spec_kernel_impl<
2581 + scalar_t, param_t, state_t>(
2582 + q.data_ptr<scalar_t>(),
2583 + k.data_ptr<scalar_t>(),
2584 + v.data_ptr<scalar_t>(),
2585 + A_log.data_ptr<param_t>(),
2586 + a.data_ptr<scalar_t>(),
2587 + dt_bias.data_ptr<scalar_t>(),
2588 + b.data_ptr<scalar_t>(),
2589 + spec_state_indices.data_ptr<int32_t>(),
2590 + num_accepted_tokens.data_ptr<int32_t>(),
2591 + cu_seqlens.data_ptr<int32_t>(),
2592 + initial_state_source.data_ptr<state_t>(),
2593 + o.data_ptr<scalar_t>(),
2594 + qk_scale_buf.data_ptr<float>(),
2595 + total_tokens,
2596 + batch_size,
2597 + spec_stride,
2598 + num_heads,
2599 + head_dim,
2600 + v_num_heads,
2601 + v_head_dim,
2602 + q_strideT,
2603 + q_strideH,
2604 + k_strideT,
2605 + k_strideH,
2606 + v_strideT,
2607 + v_strideH,
2608 + state_slot_stride,
2609 + use_qk_l2norm_in_kernel,
2610 + softplus_threshold);
2611 + });
2612 + };
2613 + if (initial_state_source.scalar_type() == at::kFloat) {
2614 + launch(float());
2615 + } else if (initial_state_source.scalar_type() == at::kBFloat16) {
2616 + launch(at::BFloat16());
2617 + } else if (initial_state_source.scalar_type() == at::kHalf) {
2618 + launch(at::Half());
2619 + } else {
2620 + TORCH_CHECK(
2621 + false,
2622 + "fused_sigmoid_gating_delta_rule_update_spec_cpu: initial_state "
2623 + "dtype must be float16, bfloat16, or float32");
2624 + }
2625 return o;
2626 }
2627
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+59 -17
@@ -37,6 +37,7 @@ CHUNK_HEAD_DIMS = [
37 CHUNK_SIZE = 64
38 CONV_DIM = 128
39 CONV_KERNEL = 4
40 +DECODE_BATCH_SIZES = [1, 3, 5]
41 PREFILL_SEQ_LENS = [
42 [1],
43 [1, 2, 3],
@@ -47,7 +48,36 @@ PREFILL_SEQ_LENS = [
48 [2 * CHUNK_SIZE - 1, 2 * CHUNK_SIZE, 2 * CHUNK_SIZE + 1],
49 [4 * CHUNK_SIZE + 17],
50 ]
50 -DECODE_BATCH_SIZES = [1, 3, 5]
51 +DECODE_DTYPE_CASES = [
52 + *[
53 + pytest.param(
54 + batch_size,
55 + num_heads,
56 + head_dims,
57 + torch.float32,
58 + id=f"fp32-b{batch_size}-h{num_heads}-d{head_dims}",
59 + )
60 + for batch_size in [1, 3, 5]
61 + for num_heads in NUM_HEADS
62 + for head_dims in HEAD_DIMS
63 + ],
64 + pytest.param(3, (2, 4), (32, 32), torch.bfloat16, id="bf16-representative"),
65 + pytest.param(3, (2, 4), (32, 32), torch.float16, id="fp16-representative"),
66 +]
67 +PREFILL_DTYPE_CASES = [
68 + *[
69 + pytest.param(
70 + num_heads,
71 + head_dims,
72 + torch.float32,
73 + id=f"fp32-h{num_heads}-d{head_dims}",
74 + )
75 + for num_heads in NUM_HEADS
76 + for head_dims in CHUNK_HEAD_DIMS
77 + ],
78 + pytest.param((2, 4), (64, 64), torch.bfloat16, id="bf16-representative"),
79 + pytest.param((2, 4), (64, 64), torch.float16, id="fp16-representative"),
80 +]
81
82
83 @functools.lru_cache(maxsize=128, typed=False)
@@ -91,6 +121,8 @@ def ref_gated_delta_rule(
121 initial_state: torch.Tensor,
122 cu_seqlens: torch.Tensor,
123 use_qk_l2norm_in_kernel: bool = False,
124 + state_dtype: torch.dtype = torch.float32,
125 + state_write_interval: int = 1,
126 ) -> tuple[torch.Tensor, torch.Tensor]:
127 g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
128 out = torch.empty_like(value)
@@ -144,6 +176,10 @@ def ref_gated_delta_rule(
176 delta = (v_t - kv_mem) * beta_t
177 state = state + delta.unsqueeze(-1) * k_t.unsqueeze(-2)
178 out_seq[:, :, token_idx] = (state * q_t.unsqueeze(-2)).sum(dim=-1)
179 + if state_dtype != torch.float32 and (
180 + (token_idx + 1) % state_write_interval == 0 or token_idx + 1 == seq_len
181 + ):
182 + state = state.to(state_dtype).float()
183
184 out[:, begin:end] = out_seq.transpose(1, 2).contiguous().to(initial_dtype)
185 final_state[seq_idx] = state.squeeze(0)
@@ -201,14 +237,15 @@ def test_fused_gdn_gating_cpu(
237
238
239 # decode path
204 -@pytest.mark.parametrize("batch_size", DECODE_BATCH_SIZES)
205 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
206 -@pytest.mark.parametrize("head_dims", HEAD_DIMS)
240 +@pytest.mark.parametrize(
241 + "batch_size,num_heads,head_dims,state_dtype", DECODE_DTYPE_CASES
242 +)
243 @torch.inference_mode()
244 def test_fused_sigmoid_gating_delta_rule_update_cpu(
245 batch_size: int,
246 num_heads: tuple[int, int],
247 head_dims: tuple[int, int],
248 + state_dtype: torch.dtype,
249 ) -> None:
250 q, k, v, a, b, A_log, dt_bias = gdn_inputs(
251 num_tokens=batch_size,
@@ -221,7 +258,7 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
258 cu_seqlens = torch.arange(batch_size + 1, dtype=torch.int32)
259 state_shape = (batch_size, num_v_heads, head_dim, v_head_dim)
260 state = tensor_cache(
224 - batch_size * num_v_heads * head_dim * v_head_dim, torch.float32
261 + batch_size * num_v_heads * head_dim * v_head_dim, state_dtype
262 ).view(state_shape)
263 state_ref = state[state_indices].transpose(-1, -2).contiguous()
264
@@ -236,6 +273,7 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
273 initial_state=state_ref,
274 cu_seqlens=cu_seqlens,
275 use_qk_l2norm_in_kernel=True,
276 + state_dtype=state_dtype,
277 )
278 out_ref = out_ref.transpose(0, 1).contiguous()
279
@@ -256,8 +294,8 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
294
295 torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
296 torch.testing.assert_close(
259 - state_out[state_indices].transpose(-1, -2),
260 - final_state_ref,
297 + state_out[state_indices].transpose(-1, -2).float(),
298 + final_state_ref.float(),
299 atol=1e-2,
300 rtol=1e-2,
301 )
@@ -265,13 +303,13 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
303
304 # prefill path
305 @pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
268 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
269 -@pytest.mark.parametrize("head_dims", CHUNK_HEAD_DIMS)
306 +@pytest.mark.parametrize("num_heads,head_dims,state_dtype", PREFILL_DTYPE_CASES)
307 @torch.inference_mode()
308 def test_chunk_gated_delta_rule_cpu(
309 seq_lens: list[int],
310 num_heads: tuple[int, int],
311 head_dims: tuple[int, int],
312 + state_dtype: torch.dtype,
313 ) -> None:
314 total_tokens = sum(seq_lens)
315 q, k, v, a, b, A_log, dt_bias = gdn_inputs(
@@ -286,7 +324,7 @@ def test_chunk_gated_delta_rule_cpu(
324 )
325 initial_state_shape = (len(seq_lens), num_v_heads, head_dim, v_head_dim)
326 initial_state = tensor_cache(
289 - len(seq_lens) * num_v_heads * head_dim * v_head_dim, torch.float32
327 + len(seq_lens) * num_v_heads * head_dim * v_head_dim, state_dtype
328 ).view(initial_state_shape)
329 initial_state_ref = initial_state.transpose(-1, -2).contiguous()
330
@@ -301,6 +339,8 @@ def test_chunk_gated_delta_rule_cpu(
339 initial_state=initial_state_ref,
340 cu_seqlens=cu_seqlens,
341 use_qk_l2norm_in_kernel=True,
342 + state_dtype=state_dtype,
343 + state_write_interval=CHUNK_SIZE,
344 )
345
346 g, beta = ref_gdn_gating(A_log, a, b, dt_bias)
@@ -320,8 +360,8 @@ def test_chunk_gated_delta_rule_cpu(
360
361 torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
362 torch.testing.assert_close(
323 - final_state.transpose(-1, -2),
324 - final_state_ref,
363 + final_state.transpose(-1, -2).float(),
364 + final_state_ref.float(),
365 atol=1e-2,
366 rtol=1e-2,
367 )
@@ -339,14 +379,16 @@ TWO_CALL_SPLITS = [
379
380
381 @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
342 -@pytest.mark.parametrize("num_heads", NUM_HEADS)
343 -@pytest.mark.parametrize("head_dims", CHUNK_HEAD_DIMS)
382 +@pytest.mark.parametrize("num_heads", [(2, 4)])
383 +@pytest.mark.parametrize("head_dims", [(64, 64)])
384 +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16, torch.float16])
385 @torch.inference_mode()
386 def test_chunk_gated_delta_rule_cpu_two_call_split(
387 total_tokens: int,
388 split: int,
389 num_heads: tuple[int, int],
390 head_dims: tuple[int, int],
391 + state_dtype: torch.dtype,
392 ) -> None:
393 """A prefill split into two calls (the second seeded with the first's
394 ``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(
407 g = g.unsqueeze(0) # [1, T, HV]
408 beta = beta.unsqueeze(0)
409
368 - zero_state = torch.zeros(1, num_v_heads, head_dim, v_head_dim, dtype=torch.float32)
410 + zero_state = torch.zeros(1, num_v_heads, head_dim, v_head_dim, dtype=state_dtype)
411
412 # Reference: whole sequence in one call, no initial state.
413 out_full, final_full = ops.chunk_gated_delta_rule_cpu(
@@ -405,7 +447,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split(
447 value=v[:, split:],
448 g=g[:, split:],
449 beta=beta[:, split:],
408 - initial_state=state1.to(torch.float32),
450 + initial_state=state1,
451 output_final_state=True,
452 cu_seqlens=torch.tensor([0, tail], dtype=torch.int32),
453 head_first=False,
@@ -416,7 +458,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split(
458 out_split = torch.cat([out1, out2], dim=1)
459
460 # State must be near-exact; output allows a looser bound for the bf16 round-trip.
419 - torch.testing.assert_close(state2, final_full, atol=1e-3, rtol=1e-3)
461 + torch.testing.assert_close(state2.float(), final_full.float(), atol=1e-3, rtol=1e-3)
462 torch.testing.assert_close(out_split, out_full, atol=2e-2, rtol=2e-2)
463
464
tests/platforms/test_cpu.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.config import CacheConfig, VllmConfig
10 +from vllm.platforms.cpu import CpuPlatform
11 +
12 +
13 +class _FakeModel:
14 + @classmethod
15 + def get_mamba_state_dtype_from_config(
16 + cls, vllm_config: VllmConfig
17 + ) -> tuple[torch.dtype, torch.dtype]:
18 + state_dtype = {
19 + "float16": torch.float16,
20 + "bfloat16": torch.bfloat16,
21 + }.get(vllm_config.cache_config.mamba_ssm_cache_dtype, torch.float32)
22 + return torch.float32, state_dtype
23 +
24 +
25 +class _FakeRegistry:
26 + @staticmethod
27 + def resolve_model_cls(
28 + architecture: str,
29 + model_config: SimpleNamespace,
30 + ) -> tuple[type[_FakeModel], str]:
31 + return _FakeModel, architecture
32 +
33 +
34 +def _cpu_config(
35 + cache_config: CacheConfig,
36 + *,
37 + model_type: str,
38 + resolved_dtype: str,
39 + architecture: str | None = None,
40 + layer_types: tuple[str, ...] = ("linear_attention",),
41 +) -> SimpleNamespace:
42 + model_config = SimpleNamespace(
43 + disable_cascade_attn=False,
44 + architecture=architecture or model_type,
45 + has_inner_state=True,
46 + use_mla=False,
47 + hf_text_config=SimpleNamespace(
48 + model_type=model_type,
49 + layer_types=layer_types,
50 + ),
51 + registry=_FakeRegistry,
52 + )
53 + config = VllmConfig(cache_config=cache_config)
54 + config.model_config = model_config
55 + cache_config.mamba_ssm_cache_dtype = resolved_dtype
56 + return config
57 +
58 +
59 +@pytest.mark.parametrize(
60 + (
61 + "model_type",
62 + "requested_dtype",
63 + "resolved_dtype",
64 + "layer_types",
65 + "expected_dtype",
66 + ),
67 + [
68 + pytest.param(
69 + "qwen3_5",
70 + "bfloat16",
71 + "bfloat16",
72 + ("linear_attention",),
73 + "bfloat16",
74 + id="gdn-explicit-bf16",
75 + ),
76 + pytest.param(
77 + "qwen3_5",
78 + "float16",
79 + "float16",
80 + ("linear_attention",),
81 + "float16",
82 + id="gdn-explicit-fp16",
83 + ),
84 + pytest.param(
85 + "qwen3_5",
86 + "auto",
87 + "bfloat16",
88 + ("linear_attention",),
89 + "bfloat16",
90 + id="gdn-model-bf16",
91 + ),
92 + pytest.param(
93 + "qwen3_5",
94 + "auto",
95 + "float16",
96 + ("linear_attention",),
97 + "float16",
98 + id="gdn-model-fp16",
99 + ),
100 + pytest.param(
101 + "nemotron_h",
102 + "auto",
103 + "bfloat16",
104 + ("mamba",),
105 + "float32",
106 + id="unsupported-model-bf16",
107 + ),
108 + pytest.param(
109 + "nemotron_h",
110 + "auto",
111 + "float16",
112 + ("mamba",),
113 + "float32",
114 + id="unsupported-model-fp16",
115 + ),
116 + pytest.param(
117 + "nemotron_h",
118 + "bfloat16",
119 + "bfloat16",
120 + ("mamba",),
121 + "float32",
122 + id="unsupported-explicit-bf16",
123 + ),
124 + pytest.param(
125 + "nemotron_h",
126 + "float16",
127 + "float16",
128 + ("mamba",),
129 + "float32",
130 + id="unsupported-explicit-fp16",
131 + ),
132 + ],
133 +)
134 +def test_cpu_accelerated_gdn_dtype_policy(
135 + monkeypatch: pytest.MonkeyPatch,
136 + model_type: str,
137 + requested_dtype: str,
138 + resolved_dtype: str,
139 + layer_types: tuple[str, ...],
140 + expected_dtype: str,
141 +) -> None:
142 + monkeypatch.setattr("torch.cpu._is_avx512_bf16_supported", lambda: True)
143 + cache_config = CacheConfig(mamba_ssm_cache_dtype=requested_dtype)
144 + config = _cpu_config(
145 + cache_config,
146 + model_type=model_type,
147 + resolved_dtype=resolved_dtype,
148 + layer_types=layer_types,
149 + )
150 +
151 + CpuPlatform.check_and_update_config(config)
152 + assert cache_config.mamba_ssm_cache_dtype == expected_dtype
vllm/platforms/cpu.py
+72 -6
@@ -31,6 +31,61 @@ else:
31 VllmConfig = None
32
33
34 +def _cpu_mamba_backend(vllm_config: VllmConfig) -> str:
35 + """Return the CPU state backend selected by the model configuration."""
36 + model_config = vllm_config.model_config
37 + if model_config is None:
38 + return "none"
39 +
40 + hf_config = getattr(model_config, "hf_text_config", None)
41 + model_type = str(getattr(hf_config, "model_type", "")).lower()
42 + architecture = str(getattr(model_config, "architecture", "")).lower()
43 + layer_types = getattr(hf_config, "layer_types", None)
44 + has_linear_attention = isinstance(layer_types, (list, tuple)) and (
45 + "linear_attention" in layer_types
46 + )
47 +
48 + try:
49 + has_inner_state = bool(model_config.has_inner_state)
50 + except (AttributeError, RuntimeError):
51 + has_inner_state = False
52 +
53 + if not (has_inner_state or has_linear_attention):
54 + return "none"
55 +
56 + fallback_backend = model_type or architecture or "unknown"
57 + if not has_linear_attention:
58 + return fallback_backend
59 +
60 + try:
61 + model_cls, _ = model_config.registry.resolve_model_cls(
62 + model_config.architecture,
63 + model_config=model_config,
64 + )
65 + except Exception:
66 + return fallback_backend
67 +
68 + try:
69 + state_dtypes = model_cls.get_mamba_state_dtype_from_config(vllm_config)
70 + except Exception:
71 + return fallback_backend
72 +
73 + if not isinstance(state_dtypes, tuple) or len(state_dtypes) != 2:
74 + return fallback_backend
75 +
76 + if vllm_config.cache_config.mamba_ssm_cache_dtype == "float16":
77 + cache_dtype = torch.float16
78 + elif vllm_config.cache_config.mamba_ssm_cache_dtype == "bfloat16":
79 + cache_dtype = torch.bfloat16
80 + else:
81 + return fallback_backend
82 +
83 + if state_dtypes[1] == cache_dtype:
84 + return "gdn"
85 +
86 + return fallback_backend
87 +
88 +
89 def get_max_threads(pid=0):
90 if hasattr(os, "sched_getaffinity"):
91 return len(os.sched_getaffinity(pid))
@@ -181,16 +236,27 @@ class CpuPlatform(Platform):
236 "otherwise the performance is not optimized."
237 )
238
184 - # Accelerated GDN (AMX tiles or AVX-512BF16 VDPBF16PS) requires
185 - # float32 SSM state.
239 + # Accelerated GDN uses AMX tiles or AVX-512BF16 VDPBF16PS.
240 if (
241 torch.cpu._is_avx512_bf16_supported()
242 and cache_config.mamba_ssm_cache_dtype != "float32"
243 ):
190 - cache_config.mamba_ssm_cache_dtype = "float32"
191 - logger.warning(
192 - "Reset SSM cache type to float32 for accelerated GDN mamba attention."
193 - )
244 + mamba_backend = _cpu_mamba_backend(vllm_config)
245 + if (
246 + cache_config.mamba_ssm_cache_dtype in ("float16", "bfloat16")
247 + and mamba_backend == "gdn"
248 + ):
249 + logger.info(
250 + "Using %s SSM state storage for the CPU accelerated GDN backend.",
251 + cache_config.mamba_ssm_cache_dtype,
252 + )
253 + else:
254 + cache_config.mamba_ssm_cache_dtype = "float32"
255 + logger.warning(
256 + "Reset SSM cache type to float32 for accelerated GDN mamba "
257 + "backend '%s'.",
258 + mamba_backend,
259 + )
260
261 # Lagecy setting
262 env_key = "VLLM_CPU_KVCACHE_SPACE"