@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 05fa8183a6

[CPU][Spec Decode] Support DFlash speculative decoding for GDN models on CPU (#46090)

Signed-off-by: guybd <guy.boudoukh@intel.com> Signed-off-by: Guy Boudoukh <guy.boudoukh@intel.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

guybd committed Jul 13, 2026 at 07:16 UTC 05fa8183a6e2c1e1c6e4643d54b9709fa29ff753

5 files changed +755 -10

csrc/cpu/sgl-kernels/fla.cpp

+255

@@ -1116,6 +1116,164 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(

1116 });

1117 }

1118

1119 +// Speculative-decode variant: processes a varlen batch where each sequence has

1120 +// ``q_len`` draft tokens, runs the recurrence sequentially over those tokens

1121 +// (inside the kernel, so one dispatch handles the whole draft block), reads the

1122 +// initial state from cache slot ``num_accepted-1`` and stores the state *after*

1123 +// token ``t`` into cache slot ``t`` (multi-slot rollback, matching the GPU

1124 +// kernel). Parallelized over (sequence, v_head); the per-sequence token loop is

1125 +// sequential as required by the recurrence.

1126 +template <typename scalar_t, typename param_t>

1127 +void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(

1128 + const scalar_t* __restrict__ q_ptr, // [T, HK, EK]

1129 + const scalar_t* __restrict__ k_ptr, // [T, HK, EK]

1130 + const scalar_t* __restrict__ v_ptr, // [T, HV, EV]

1131 + const param_t* __restrict__ A_log_ptr,

1132 + const scalar_t* __restrict__ a_ptr, // [T, HV]

1133 + const scalar_t* __restrict__ dt_bias_ptr,

1134 + const scalar_t* __restrict__ b_ptr, // [T, HV]

1135 + const int32_t* __restrict__ spec_indices_ptr, // [N, S]

1136 + const int32_t* __restrict__ num_accepted_ptr, // [N]

1137 + const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1]

1138 + float* __restrict__ state_ptr,

1139 + scalar_t* __restrict__ o_ptr, // [T, HV, EV]

1140 + float* __restrict__ qk_scale_buf, // [2, T, HK]

1141 + int64_t total_tokens,

1142 + int64_t batch_size,

1143 + int64_t spec_stride,

1144 + int64_t num_heads,

1145 + int64_t head_dim,

1146 + int64_t v_num_heads,

1147 + int64_t v_head_dim,

1148 + int64_t q_strideT,

1149 + int64_t q_strideH,

1150 + int64_t k_strideT,

1151 + int64_t k_strideH,

1152 + int64_t v_strideT,

1153 + int64_t v_strideH,

1154 + int64_t state_slot_stride,

1155 + bool use_qk_l2norm_in_kernel,

1156 + double softplus_threshold) {

1157 + using bVec = at::vec::Vectorized<scalar_t>;

1158 + using fVec = at::vec::Vectorized<float>;

1159 + constexpr int64_t VecSize = bVec::size();

1160 + constexpr int64_t fVecSize = fVec::size();

1161 + int64_t group_size = v_num_heads / num_heads;

1162 + double scale = 1 / std::sqrt((double)head_dim);

1163 + fVec scale_vec = fVec((float)scale);

1164 +

1165 + if (use_qk_l2norm_in_kernel) {

1166 + float eps = 1e-5f;

1167 + at::parallel_for(0, total_tokens * num_heads, 0, [&](int64_t begin, int64_t end) {

1168 + for (int64_t i = begin; i < end; ++i) {

1169 + int64_t ti = i / num_heads;

1170 + int64_t ni = i % num_heads;

1171 + const scalar_t* qp = q_ptr + ti * q_strideT + ni * q_strideH;

1172 + const scalar_t* kp = k_ptr + ti * k_strideT + ni * k_strideH;

1173 + float sq = 0.f, sk = 0.f;

1174 + for (int64_t d = 0; d < head_dim; ++d) {

1175 + float qv = (float)qp[d];

1176 + sq += qv * qv;

1177 + float kv = (float)kp[d];

1178 + sk += kv * kv;

1179 + }

1180 + qk_scale_buf[ti * num_heads + ni] = 1.f / std::sqrt(sq + eps);

1181 + qk_scale_buf[total_tokens * num_heads + ti * num_heads + ni] = 1.f / std::sqrt(sk + eps);

1182 + }

1183 + });

1184 + }

1185 +

1186 + at::parallel_for(0, batch_size * v_num_heads, 0, [&](int64_t begin, int64_t end) {

1187 + for (int64_t idx = begin; idx < end; ++idx) {

1188 + int64_t bi = idx / v_num_heads;

1189 + int64_t ni = idx % v_num_heads;

1190 + int64_t kh = ni / group_size;

1191 + int64_t q_start = cu_seqlens_ptr[bi];

1192 + int64_t q_len = cu_seqlens_ptr[bi + 1] - q_start;

1193 + if (q_len <= 0) {

1194 + continue;

1195 + }

1196 + int64_t acc = (int64_t)num_accepted_ptr[bi];

1197 + // Clamp acc-1 to >=0: when num_accepted is 0 the unclamped index reads

1198 + // out of bounds and yields an arbitrary prev_slot used to index the SSM

1199 + // state. Mirrors the GPU guard tl.maximum(num_accepted - 1, 0).

1200 + int64_t prev_slot =

1201 + (int64_t)spec_indices_ptr[bi * spec_stride + (acc > 0 ? acc - 1 : 0)];

1202 + for (int64_t t = 0; t < q_len; ++t) {

1203 + int64_t cur_slot = (int64_t)spec_indices_ptr[bi * spec_stride + t];

1204 + int64_t token = q_start + t;

1205 + const float* src = state_ptr + prev_slot * state_slot_stride + ni * head_dim * v_head_dim;

1206 + float* dst = state_ptr + cur_slot * state_slot_stride + ni * head_dim * v_head_dim;

1207 + float g_val = -std::exp((float)A_log_ptr[ni]) *

1208 + softplus((float)a_ptr[token * v_num_heads + ni] + (float)dt_bias_ptr[ni], softplus_threshold);

1209 + float g_val_exp = std::exp(g_val);

1210 + fVec g_val_exp_vec = fVec(g_val_exp);

1211 + float beta_val = 1.f / (1.f + std::exp(-(float)b_ptr[token * v_num_heads + ni]));

1212 + fVec beta_vec = fVec(beta_val);

1213 + int64_t q_offset = token * q_strideT + kh * q_strideH;

1214 + int64_t k_offset = token * k_strideT + kh * k_strideH;

1215 + float q_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[token * num_heads + kh] : 1.f;

1216 + float k_scale =

1217 + use_qk_l2norm_in_kernel ? qk_scale_buf[total_tokens * num_heads + token * num_heads + kh] : 1.f;

1218 + int64_t v_offset = token * v_strideT + ni * v_strideH;

1219 + int64_t o_offset = (token * v_num_heads + ni) * v_head_dim;

1220 + int64_t dvi = 0;

1221 + for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {

1222 + fVec kv_mem_vec0 = fVec(0.f);

1223 + fVec kv_mem_vec1 = fVec(0.f);

1224 + for (int di = 0; di < head_dim; ++di) {

1225 + fVec k_val_vec = fVec((float)k_ptr[k_offset + di] * k_scale);

1226 + fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi);

1227 + fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize);

1228 + kv_mem_vec0 = kv_mem_vec0 + sv0 * g_val_exp_vec * k_val_vec;

1229 + kv_mem_vec1 = kv_mem_vec1 + sv1 * g_val_exp_vec * k_val_vec;

1230 + }

1231 + bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi);

1232 + fVec v_vec0, v_vec1;

1233 + std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec);

1234 + fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec;

1235 + fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec;

1236 + fVec o_vec0 = fVec(0.f);

1237 + fVec o_vec1 = fVec(0.f);

1238 + for (int di = 0; di < head_dim; ++di) {

1239 + fVec q_vec = fVec((float)q_ptr[q_offset + di] * q_scale);

1240 + fVec k_vec = fVec((float)k_ptr[k_offset + di] * k_scale);

1241 + fVec sv0 = fVec::loadu(src + di * v_head_dim + dvi);

1242 + fVec sv1 = fVec::loadu(src + di * v_head_dim + dvi + fVecSize);

1243 + sv0 = sv0 * g_val_exp_vec + k_vec * dt_vec0;

1244 + sv1 = sv1 * g_val_exp_vec + k_vec * dt_vec1;

1245 + o_vec0 = o_vec0 + sv0 * q_vec * scale_vec;

1246 + o_vec1 = o_vec1 + sv1 * q_vec * scale_vec;

1247 + sv0.store(dst + di * v_head_dim + dvi);

1248 + sv1.store(dst + di * v_head_dim + dvi + fVecSize);

1249 + }

1250 + bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1);

1251 + o_vec.store(o_ptr + o_offset + dvi);

1252 + }

1253 + for (; dvi < v_head_dim; ++dvi) {

1254 + float kv_mem_val = 0.f;

1255 + for (int di = 0; di < head_dim; ++di) {

1256 + float k_val = (float)k_ptr[k_offset + di] * k_scale;

1257 + kv_mem_val += src[di * v_head_dim + dvi] * g_val_exp * k_val;

1258 + }

1259 + float v_val = (float)v_ptr[v_offset + dvi];

1260 + float dt_val = (v_val - kv_mem_val) * beta_val;

1261 + float o_val = 0.f;

1262 + for (int di = 0; di < head_dim; ++di) {

1263 + float q_val = (float)q_ptr[q_offset + di] * q_scale;

1264 + float k_val = (float)k_ptr[k_offset + di] * k_scale;

1265 + float ns = src[di * v_head_dim + dvi] * g_val_exp + k_val * dt_val;

1266 + dst[di * v_head_dim + dvi] = ns;

1267 + o_val += ns * q_val * scale;

1268 + }

1269 + o_ptr[o_offset + dvi] = (scalar_t)o_val;

1270 + }

1271 + prev_slot = cur_slot;

1272 + }

1273 + }

1274 + });

1275 +}

1276 +

1277 template <typename scalar_t>

1278 void fused_gdn_gating_kernel_impl(

1279 float* __restrict__ A_log,

@@ -1500,6 +1658,103 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(

1658 return core_attn_out;

1659 }

1660

1661 +// Speculative-decode update (multi-token, multi-slot rollback).

1662 +// q: [T, HK, EK] k: [T, HK, EK] v: [T, HV, EV]

1663 +// a: [T, HV] b: [T, HV]

1664 +// initial_state_source: [N_slots, HV, EK, EV] FP32 (updated in place)

1665 +// spec_state_indices: [batch, S] INT32 (S = num_spec + 1)

1666 +// num_accepted_tokens: [batch] INT32

1667 +// cu_seqlens: [batch + 1] INT32

1668 +// Returns output: [T, HV, EV]

1669 +at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(

1670 + const at::Tensor& A_log,

1671 + const at::Tensor& dt_bias,

1672 + const at::Tensor& q,

1673 + const at::Tensor& k,

1674 + const at::Tensor& v,

1675 + const at::Tensor& a,

1676 + const at::Tensor& b,

1677 + at::Tensor& initial_state_source,

1678 + const at::Tensor& spec_state_indices,

1679 + const at::Tensor& num_accepted_tokens,

1680 + const at::Tensor& cu_seqlens,

1681 + bool use_qk_l2norm_in_kernel,

1682 + double softplus_beta = 1.0,

1683 + double softplus_threshold = 20.0) {

1684 + CHECK_DIM(3, q);

1685 + CHECK_DIM(3, v);

1686 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);

1687 + int64_t total_tokens = q.size(0);

1688 + int64_t num_heads = q.size(1);

1689 + int64_t head_dim = q.size(2);

1690 + int64_t v_num_heads = v.size(1);

1691 + int64_t v_head_dim = v.size(2);

1692 + int64_t batch_size = cu_seqlens.size(0) - 1;

1693 + int64_t spec_stride = spec_state_indices.stride(0);

1694 + CHECK_INPUT_SHAPE_DTYPE<true>(k, {total_tokens, num_heads, head_dim}, q.scalar_type());

1695 + CHECK_INPUT_SHAPE_DTYPE<true>(v, {total_tokens, v_num_heads, v_head_dim}, q.scalar_type());

1696 + CHECK_INPUT_SHAPE_DTYPE<true>(a, {total_tokens, v_num_heads}, q.scalar_type());

1697 + CHECK_INPUT_SHAPE_DTYPE<true>(b, {total_tokens, v_num_heads}, q.scalar_type());

1698 + CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type());

1699 + CHECK_INPUT_SHAPE_DTYPE<true>(num_accepted_tokens, {batch_size}, at::kInt);

1700 + CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt);

1701 + CHECK_EQ(v_num_heads % num_heads, 0);

1702 + TORCH_CHECK(A_log.sizes() == at::IntArrayRef({v_num_heads}));

1703 + CHECK_INPUT_SHAPE_DTYPE<true>(

1704 + initial_state_source,

1705 + {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim},

1706 + at::kFloat);

1707 + TORCH_CHECK(initial_state_source.size(0) >= batch_size,

1708 + "initial_state_source capacity too small: size(0)=",

1709 + initial_state_source.size(0), ", batch_size=", batch_size);

1710 +

1711 + int64_t q_strideT = q.stride(0);

1712 + int64_t q_strideH = q.stride(1);

1713 + int64_t k_strideT = k.stride(0);

1714 + int64_t k_strideH = k.stride(1);

1715 + int64_t v_strideT = v.stride(0);

1716 + int64_t v_strideH = v.stride(1);

1717 + int64_t state_slot_stride = initial_state_source.stride(0);

1718 +

1719 + at::Tensor o = at::empty({total_tokens, v_num_heads, v_head_dim}, q.options());

1720 + at::Tensor qk_scale_buf = at::empty({2, total_tokens, num_heads}, at::kFloat);

1721 +

1722 + CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(

1723 + q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_spec_kernel_impl", [&] {

1724 + fused_sigmoid_gating_delta_rule_update_spec_kernel_impl<scalar_t, param_t>(

1725 + q.data_ptr<scalar_t>(),

1726 + k.data_ptr<scalar_t>(),

1727 + v.data_ptr<scalar_t>(),

1728 + A_log.data_ptr<param_t>(),

1729 + a.data_ptr<scalar_t>(),

1730 + dt_bias.data_ptr<scalar_t>(),

1731 + b.data_ptr<scalar_t>(),

1732 + spec_state_indices.data_ptr<int32_t>(),

1733 + num_accepted_tokens.data_ptr<int32_t>(),

1734 + cu_seqlens.data_ptr<int32_t>(),

1735 + initial_state_source.data_ptr<float>(),

1736 + o.data_ptr<scalar_t>(),

1737 + qk_scale_buf.data_ptr<float>(),

1738 + total_tokens,

1739 + batch_size,

1740 + spec_stride,

1741 + num_heads,

1742 + head_dim,

1743 + v_num_heads,

1744 + v_head_dim,

1745 + q_strideT,

1746 + q_strideH,

1747 + k_strideT,

1748 + k_strideH,

1749 + v_strideT,

1750 + v_strideH,

1751 + state_slot_stride,

1752 + use_qk_l2norm_in_kernel,

1753 + softplus_threshold);

1754 + });

1755 + return o;

1756 +}

1757 +

1758 // A_log: [num_v_heads]

1759 // a: [batch, num_v_heads]

1760 // b: [batch, num_v_heads]

csrc/cpu/torch_bindings.cpp

+17

@@ -120,6 +120,14 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(

120 bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0,

121 double softplus_threshold = 20.0);

122

123 +at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(

124 + const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q,

125 + const at::Tensor& k, const at::Tensor& v, const at::Tensor& a,

126 + const at::Tensor& b, at::Tensor& initial_state_source,

127 + const at::Tensor& spec_state_indices, const at::Tensor& num_accepted_tokens,

128 + const at::Tensor& cu_seqlens, bool use_qk_l2norm_in_kernel,

129 + double softplus_beta = 1.0, double softplus_threshold = 20.0);

130 +

131 std::tuple<at::Tensor, at::Tensor> fused_gdn_gating_cpu(

132 const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b,

133 const at::Tensor& dt_bias);

@@ -508,6 +516,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

516 "softplus_threshold=20.0) -> Tensor");

517 ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU,

518 &fused_sigmoid_gating_delta_rule_update_cpu);

519 + ops.def(

520 + "fused_sigmoid_gating_delta_rule_update_spec_cpu(Tensor A_log, Tensor "

521 + "dt_bias, Tensor q, Tensor k, Tensor v, Tensor a, Tensor b, "

522 + "Tensor(a!) initial_state_source, Tensor spec_state_indices, "

523 + "Tensor num_accepted_tokens, Tensor cu_seqlens, bool "

524 + "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float "

525 + "softplus_threshold=20.0) -> Tensor");

526 + ops.impl("fused_sigmoid_gating_delta_rule_update_spec_cpu", torch::kCPU,

527 + &fused_sigmoid_gating_delta_rule_update_spec_cpu);

528 ops.def(

529 "fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) "

530 "-> (Tensor, Tensor)");

vllm/_custom_ops.py

+34

@@ -3252,6 +3252,40 @@ def fused_sigmoid_gating_delta_rule_update_cpu(

3252 )

3253

3254

3255 +def fused_sigmoid_gating_delta_rule_update_spec_cpu(

3256 + A_log: torch.Tensor,

3257 + dt_bias: torch.Tensor,

3258 + q: torch.Tensor,

3259 + k: torch.Tensor,

3260 + v: torch.Tensor,

3261 + a: torch.Tensor,

3262 + b: torch.Tensor,

3263 + initial_state_source: torch.Tensor,

3264 + spec_state_indices: torch.Tensor,

3265 + num_accepted_tokens: torch.Tensor,

3266 + cu_seqlens: torch.Tensor,

3267 + use_qk_l2norm_in_kernel: bool,

3268 + softplus_beta: float = 1.0,

3269 + softplus_threshold: float = 20.0,

3270 +) -> torch.Tensor:

3271 + return torch.ops._C.fused_sigmoid_gating_delta_rule_update_spec_cpu(

3272 + A_log,

3273 + dt_bias,

3274 + q,

3275 + k,

3276 + v,

3277 + a,

3278 + b,

3279 + initial_state_source,

3280 + spec_state_indices,

3281 + num_accepted_tokens,

3282 + cu_seqlens,

3283 + use_qk_l2norm_in_kernel,

3284 + softplus_beta,

3285 + softplus_threshold,

3286 + )

3287 +

3288 +

3289 def fused_gdn_gating_cpu(

3290 A_log: torch.Tensor,

3291 a: torch.Tensor,

vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py

+442 -9

@@ -4,6 +4,7 @@

4 from __future__ import annotations

6 import torch

7 +import torch.nn.functional as F

9 import vllm._custom_ops as ops

10 from vllm.forward_context import ForwardContext, get_forward_context

@@ -46,11 +47,52 @@ def cpu_gdn_attention_core(

47 if attn_metadata_i.num_actual_tokens == 0:

48 return

50 + assert mixed_qkv.dtype == torch.bfloat16, "CPU GDN attention requires BF16."

51 +

52 + # The conv-state cache temporal dim is ``width - 1`` when speculative

53 + # decoding is disabled, and ``width - 1 + num_spec`` when it is enabled

54 + # (vLLM allocates a wider rolling buffer to support draft-token rollback).

55 + width = layer.conv1d.weight.size(-1)

56 + conv_cache = layer.kv_cache[0]

57 + if is_conv_state_dim_first():

58 + # DS layout: (num_slots, dim, state_len)

59 + state_len = conv_cache.size(-1)

60 + else:

61 + # SD layout: (num_slots, state_len, dim)

62 + state_len = conv_cache.size(-2)

63 +

64 + spec_decode_cache = state_len > (width - 1)

65 +

66 + if not spec_decode_cache:

67 + # ===== Fast path: speculative decoding NOT configured =====

68 + # conv-state is exactly (..., dim, width-1); use the original AMX /

69 + # torch implementation unchanged.

70 + _cpu_gdn_attention_nonspec(

71 + layer, attn_metadata_i, mixed_qkv, b, a, core_attn_out

72 + )

73 + return

74 +

75 + # ===== Speculative-decoding-aware path (wide rolling conv buffer) =====

76 + _cpu_gdn_attention_spec_aware(

77 + layer, attn_metadata_i, mixed_qkv, b, a, core_attn_out, width, state_len

78 + )

79 +

80 +

81 +# ---------------------------------------------------------------------------

82 +# Non-speculative implementation: conv-state is exactly width-1.

83 +# ---------------------------------------------------------------------------

84 +def _cpu_gdn_attention_nonspec(

85 + layer,

86 + attn_metadata_i: GDNAttentionMetadata,

87 + mixed_qkv: torch.Tensor,

88 + b: torch.Tensor,

89 + a: torch.Tensor,

90 + core_attn_out: torch.Tensor,

91 +) -> None:

92 assert (

93 attn_metadata_i.spec_sequence_masks is None

94 and attn_metadata_i.num_accepted_tokens is None

52 - ), "speculative decode not supported in CPU GDN attention."

53 - assert mixed_qkv.dtype == torch.bfloat16, "CPU GDN attention requires BF16."

95 + ), "speculative decode not supported without a wide conv-state cache."

97 state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor

98 query_start_loc = attn_metadata_i.non_spec_query_start_loc

@@ -79,12 +121,7 @@ def cpu_gdn_attention_core(

121 b = b.contiguous()

122

123 num_allocated_slots, head_num, v_dim, k_dim = ssm_state.size()

82 - ssm_state = ssm_state.view(

83 - num_allocated_slots,

84 - head_num,

85 - k_dim,

86 - v_dim,

87 - )

124 + ssm_state = ssm_state.view(num_allocated_slots, head_num, k_dim, v_dim)

125

126 num_decodes = attn_metadata_i.num_decodes

127 num_decode_tokens = attn_metadata_i.num_decode_tokens

@@ -109,7 +146,6 @@ def cpu_gdn_attention_core(

146 )

147 else:

148 decode_conv_state = conv_state[decode_state_indices].contiguous()

112 -

149 decode_mixed_qkv = causal_conv1d_update_torch(

150 # [B, dim] -> [B, dim, 1]

151 x=decode_mixed_qkv.unsqueeze(-1),

@@ -207,6 +243,403 @@ def cpu_gdn_attention_core(

243 core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)

244

245

246 +# ---------------------------------------------------------------------------

247 +# Speculative-decoding-aware implementation (wide rolling conv buffer).

248 +#

249 +# Mirrors the GPU `_forward_core` spec path using torch + the existing CPU

250 +# kernels. The conv-state is a per-sequence rolling buffer of length

251 +# ``state_len = width - 1 + num_spec`` (single cache slot, column 0 of the

252 +# spec block table), and the SSM recurrent state is stored multi-slot: the

253 +# state *after* draft token ``t`` lives in block-table column ``t``, so the

254 +# next step can resume from the slot of the last accepted token.

255 +# ---------------------------------------------------------------------------

256 +def _conv_buffer_view(layer) -> torch.Tensor:

257 + """Return the conv-state cache as (num_slots, dim, state_len)."""

258 + conv_cache = layer.kv_cache[0]

259 + if is_conv_state_dim_first():

260 + return conv_cache # (num_slots, dim, state_len)

261 + return conv_cache.transpose(-1, -2) # view -> (num_slots, dim, state_len)

262 +

263 +

264 +def _ssm_state_view(layer) -> torch.Tensor:

265 + ssm_state = layer.kv_cache[1]

266 + num_slots, head_num, v_dim, k_dim = ssm_state.size()

267 + return ssm_state.view(num_slots, head_num, k_dim, v_dim)

268 +

269 +

270 +def _unpacked_conv_weight(layer) -> torch.Tensor:

271 + """Return the plain (dim, width) conv weight.

272 +

273 + On AMX the conv1d weight is VNNI-packed in place at load time (only usable

274 + by the AMX C++ kernel), so the torch spec-decode path relies on the

275 + un-packed copy stashed by ``dispatch_cpu_unquantized_gemm``.

276 + """

277 + w = getattr(layer.conv1d, "_cpu_unpacked_conv_weight", None)

278 + if w is not None:

279 + return w

280 + w = layer.conv1d.weight

281 + if w.dim() == 3:

282 + return w.view(w.size(0), w.size(2))

283 + return w

284 +

285 +

286 +def _cpu_gdn_attention_spec_aware(

287 + layer,

288 + attn_metadata_i: GDNAttentionMetadata,

289 + mixed_qkv: torch.Tensor,

290 + b: torch.Tensor,

291 + a: torch.Tensor,

292 + core_attn_out: torch.Tensor,

293 + width: int,

294 + state_len: int,

295 +) -> None:

296 + mixed_qkv = mixed_qkv.contiguous()

297 + a = a.contiguous()

298 + b = b.contiguous()

299 +

300 + spec_sequence_masks = attn_metadata_i.spec_sequence_masks

301 + conv_buf = _conv_buffer_view(layer) # (num_slots, dim, state_len)

302 + ssm_state = _ssm_state_view(layer)

303 +

304 + if spec_sequence_masks is None:

305 + # No spec sequences in this batch (e.g. the prompt prefill step while

306 + # speculative decoding is configured). Process as prefill/decode using

307 + # torch conv (which only touches the first ``width-1`` columns of the

308 + # wide buffer, leaving the rolling history untouched).

309 + _spec_aware_nonspec(

310 + layer,

311 + attn_metadata_i,

312 + mixed_qkv,

313 + b,

314 + a,

315 + core_attn_out,

316 + conv_buf,

317 + ssm_state,

318 + width,

319 + )

320 + return

321 +

322 + # Split spec vs non-spec tokens.

323 + spec_token_indx = attn_metadata_i.spec_token_indx

324 + non_spec_token_indx = attn_metadata_i.non_spec_token_indx

325 + num_prefills = attn_metadata_i.num_prefills

326 + num_decodes = attn_metadata_i.num_decodes

327 +

328 + if num_prefills == 0 and num_decodes == 0:

329 + mixed_qkv_spec = mixed_qkv

330 + b_spec = b

331 + a_spec = a

332 + spec_out_indx = None

333 + else:

334 + assert spec_token_indx is not None

335 + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)

336 + b_spec = b.index_select(0, spec_token_indx)

337 + a_spec = a.index_select(0, spec_token_indx)

338 + spec_out_indx = spec_token_indx

339 +

340 + spec_out = _spec_forward(

341 + layer,

342 + attn_metadata_i,

343 + mixed_qkv_spec,

344 + b_spec,

345 + a_spec,

346 + conv_buf,

347 + ssm_state,

348 + width,

349 + state_len,

350 + )

351 +

352 + # Process the (rare) non-spec subset (prefill) tokens, if any.

353 + nonspec_out = None

354 + if (num_prefills > 0 or num_decodes > 0) and non_spec_token_indx is not None:

355 + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx)

356 + b_ns = b.index_select(0, non_spec_token_indx)

357 + a_ns = a.index_select(0, non_spec_token_indx)

358 + nonspec_out = _spec_aware_nonspec_subset(

359 + layer, attn_metadata_i, mixed_qkv_ns, b_ns, a_ns, conv_buf, ssm_state, width

360 + )

361 +

362 + # Scatter outputs back to their token positions.

363 + if spec_out_indx is None:

364 + core_attn_out[: spec_out.size(0)] = spec_out

365 + else:

366 + core_attn_out.index_copy_(0, spec_out_indx, spec_out)

367 + if nonspec_out is not None:

368 + core_attn_out.index_copy_(0, non_spec_token_indx, nonspec_out)

369 +

370 +

371 +def _spec_forward(

372 + layer,

373 + attn_metadata_i: GDNAttentionMetadata,

374 + mixed_qkv_spec: torch.Tensor,

375 + b_spec: torch.Tensor,

376 + a_spec: torch.Tensor,

377 + conv_buf: torch.Tensor,

378 + ssm_state: torch.Tensor,

379 + width: int,

380 + state_len: int,

381 +) -> torch.Tensor:

382 + """Run the GDN core for the multi-query (speculative) tokens.

383 +

384 + Returns the core attention output for the spec tokens, in the same token

385 + order as ``mixed_qkv_spec`` (i.e. ordered by ``spec_query_start_loc``).

386 + """

387 + num_spec_decodes = attn_metadata_i.num_spec_decodes

388 + spec_state_indices = attn_metadata_i.spec_state_indices_tensor

389 + spec_qsl = attn_metadata_i.spec_query_start_loc

390 + num_accepted = attn_metadata_i.num_accepted_tokens

391 + assert spec_state_indices is not None

392 + assert spec_qsl is not None

393 + assert num_accepted is not None

394 +

395 + spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1].to("cpu", torch.int64)

396 + num_acc_cpu = num_accepted[:num_spec_decodes].to("cpu", torch.int64)

397 + seq_starts = spec_qsl_cpu[:-1]

398 + seq_lens = spec_qsl_cpu[1:] - spec_qsl_cpu[:-1]

399 +

400 + # ---- 1. Convolution (per-sequence rolling buffer) ----

401 + w2d = _unpacked_conv_weight(layer) # (dim, width)

402 + dim = w2d.size(0)

403 + w = w2d.unsqueeze(1) # (dim, 1, width) for F.conv1d depthwise

404 + bias = layer.conv1d.bias

405 + silu = layer.activation == "silu"

406 +

407 + conv_out = torch.empty_like(mixed_qkv_spec)

408 + col0 = spec_state_indices[:, 0].to("cpu", torch.int64)

409 + for i in range(num_spec_decodes):

410 + q_i = int(seq_lens[i].item())

411 + if q_i == 0:

412 + continue

413 + start = int(seq_starts[i].item())

414 + slot0 = int(col0[i].item())

415 + a_prev = int(num_acc_cpu[i].item())

416 + offset = a_prev - 1

417 + B = conv_buf[slot0] # (dim, state_len)

418 + x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype)

419 + prior = B[:, offset : offset + (width - 1)]

420 + conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) # (1, dim, w-1+q)

421 + out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i)

422 + if silu:

423 + out = F.silu(out)

424 + conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype)

425 + # Roll the buffer: drop ``a_prev`` from the front, append the new

426 + # draft tokens, keep total length == state_len.

427 + keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)]

428 + new_B = torch.cat([keep, x_seq], dim=-1)

429 + B.copy_(new_B)

430 +

431 + # ---- 2. Recurrent (multi-slot SSM state) ----

432 + # Single fused kernel call: it runs the recurrence over each sequence's

433 + # draft tokens internally, resumes from slot ``num_accepted-1`` and stores

434 + # the state after token ``t`` into slot ``t`` (rollback for the next step).

435 + query, key, value = layer.rearrange_mixed_qkv(conv_out)

436 + query = query.squeeze(0).contiguous()

437 + key = key.squeeze(0).contiguous()

438 + value = value.squeeze(0).contiguous()

439 + spec_idx = spec_state_indices[:num_spec_decodes].to(torch.int32).contiguous()

440 + num_acc = num_accepted[:num_spec_decodes].to(torch.int32).contiguous()

441 + cu = spec_qsl[: num_spec_decodes + 1].to(torch.int32).contiguous()

442 + out_spec = ops.fused_sigmoid_gating_delta_rule_update_spec_cpu(

443 + A_log=layer.A_log,

444 + dt_bias=layer.dt_bias,

445 + q=query,

446 + k=key,

447 + v=value,

448 + a=a_spec.contiguous(),

449 + b=b_spec.contiguous(),

450 + initial_state_source=ssm_state,

451 + spec_state_indices=spec_idx,

452 + num_accepted_tokens=num_acc,

453 + cu_seqlens=cu,

454 + use_qk_l2norm_in_kernel=True,

455 + )

456 + return out_spec

457 +

458 +

459 +def core_attn_out_like(layer, mixed_qkv_spec: torch.Tensor) -> torch.Tensor:

460 + num_tokens = mixed_qkv_spec.size(0)

461 + return torch.zeros(

462 + (num_tokens, layer.num_v_heads // layer.tp_size, layer.head_v_dim),

463 + dtype=mixed_qkv_spec.dtype,

464 + device=mixed_qkv_spec.device,

465 + )

466 +

467 +

468 +def _spec_aware_nonspec(

469 + layer,

470 + attn_metadata_i: GDNAttentionMetadata,

471 + mixed_qkv: torch.Tensor,

472 + b: torch.Tensor,

473 + a: torch.Tensor,

474 + core_attn_out: torch.Tensor,

475 + conv_buf: torch.Tensor,

476 + ssm_state: torch.Tensor,

477 + width: int,

478 +) -> None:

479 + """Non-spec prefill/decode with a wide conv buffer (torch path)."""

480 + state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor

481 + query_start_loc = attn_metadata_i.non_spec_query_start_loc

482 + assert state_indices_tensor is not None

483 + assert query_start_loc is not None

484 +

485 + conv_weights = _unpacked_conv_weight(layer)

486 +

487 + num_decodes = attn_metadata_i.num_decodes

488 + num_decode_tokens = attn_metadata_i.num_decode_tokens

489 + num_prefills = attn_metadata_i.num_prefills

490 + num_prefill_tokens = attn_metadata_i.num_prefill_tokens

491 +

492 + if num_decodes > 0:

493 + decode_mixed_qkv = mixed_qkv[:num_decode_tokens]

494 + decode_b = b[:num_decode_tokens]

495 + decode_a = a[:num_decode_tokens]

496 + decode_state_indices = state_indices_tensor[:num_decodes]

497 + # Only the first ``width-1`` columns hold the real conv state.

498 + decode_conv_state = conv_buf[decode_state_indices][

499 + :, :, : width - 1

500 + ].contiguous()

501 + decode_mixed_qkv = causal_conv1d_update_torch(

502 + x=decode_mixed_qkv.unsqueeze(-1),

503 + conv_state=decode_conv_state,

504 + weight=conv_weights,

505 + bias=layer.conv1d.bias,

506 + activation=layer.activation,

507 + ).squeeze(-1)

508 + conv_buf[decode_state_indices, :, : width - 1] = decode_conv_state

509 +

510 + query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)

511 + # rearrange_mixed_qkv can return views whose last dim is not

512 + # contiguous; the fused CPU kernel requires a contiguous last dim.

513 + query = query.contiguous()

514 + key = key.contiguous()

515 + value = value.contiguous()

516 + attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu(

517 + A_log=layer.A_log,

518 + dt_bias=layer.dt_bias,

519 + q=query,

520 + k=key,

521 + v=value,

522 + a=decode_a.contiguous(),

523 + b=decode_b.contiguous(),

524 + initial_state_source=ssm_state,

525 + initial_state_indices=decode_state_indices,

526 + cu_seqlens=query_start_loc[: num_decodes + 1],

527 + use_qk_l2norm_in_kernel=True,

528 + )

529 + core_attn_out[:num_decode_tokens] = attn_out.squeeze(1)

530 +

531 + if num_prefills > 0:

532 + has_initial_state = attn_metadata_i.has_initial_state

533 + assert has_initial_state is not None

534 + prefill_token_start = num_decode_tokens

535 + prefill_token_end = prefill_token_start + num_prefill_tokens

536 + prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end]

537 + prefill_b = b[prefill_token_start:prefill_token_end]

538 + prefill_a = a[prefill_token_start:prefill_token_end]

539 + prefill_state_indices = state_indices_tensor[

540 + num_decodes : num_decodes + num_prefills

541 + ]

542 + prefill_query_start_loc = (

543 + query_start_loc[num_decodes : num_decodes + num_prefills + 1]

544 + - num_decode_tokens

545 + )

546 + prefill_has_initial_state = has_initial_state[

547 + num_decodes : num_decodes + num_prefills

548 + ]

549 + # ``causal_conv1d_torch`` only touches columns [:width-1] of the buffer.

550 + prefill_mixed_qkv = causal_conv1d_torch(

551 + x=prefill_mixed_qkv.transpose(0, 1),

552 + weight=conv_weights,

553 + bias=layer.conv1d.bias,

554 + conv_states=conv_buf,

555 + query_start_loc=prefill_query_start_loc,

556 + cache_indices=prefill_state_indices,

557 + has_initial_state=prefill_has_initial_state,

558 + activation=layer.activation,

559 + ).transpose(0, 1)

560 +

561 + query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv)

562 + g, beta = ops.fused_gdn_gating_cpu(

563 + A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias

564 + )

565 + initial_state = ssm_state[prefill_state_indices]

566 + initial_state[~prefill_has_initial_state, ...] = 0

567 + attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu(

568 + query=query,

569 + key=key,

570 + value=value,

571 + g=g,

572 + beta=beta,

573 + initial_state=initial_state,

574 + output_final_state=True,

575 + cu_seqlens=prefill_query_start_loc,

576 + head_first=False,

577 + use_qk_l2norm_in_kernel=True,

578 + )

579 + ssm_state[prefill_state_indices] = last_recurrent_state.to(

580 + ssm_state.dtype, copy=False

581 + )

582 + core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)

583 +

584 +

585 +def _spec_aware_nonspec_subset(

586 + layer,

587 + attn_metadata_i: GDNAttentionMetadata,

588 + mixed_qkv: torch.Tensor,

589 + b: torch.Tensor,

590 + a: torch.Tensor,

591 + conv_buf: torch.Tensor,

592 + ssm_state: torch.Tensor,

593 + width: int,

594 +) -> torch.Tensor:

595 + """Process non-spec (prefill) tokens that coexist with spec sequences.

596 +

597 + Returns outputs ordered like ``non_spec_token_indx``.

598 + """

599 + out = core_attn_out_like(layer, mixed_qkv)

600 + has_initial_state = attn_metadata_i.has_initial_state

601 + prefill_state_indices = attn_metadata_i.prefill_state_indices

602 + prefill_qsl = attn_metadata_i.prefill_query_start_loc

603 + assert prefill_state_indices is not None and prefill_qsl is not None

604 + assert has_initial_state is not None

605 +

606 + conv_weights = _unpacked_conv_weight(layer)

607 + conv_out = causal_conv1d_torch(

608 + x=mixed_qkv.transpose(0, 1),

609 + weight=conv_weights,

610 + bias=layer.conv1d.bias,

611 + conv_states=conv_buf,

612 + query_start_loc=prefill_qsl,

613 + cache_indices=prefill_state_indices,

614 + has_initial_state=has_initial_state,

615 + activation=layer.activation,

616 + ).transpose(0, 1)

617 +

618 + query, key, value = layer.rearrange_mixed_qkv(conv_out)

619 + g, beta = ops.fused_gdn_gating_cpu(

620 + A_log=layer.A_log, a=a, b=b, dt_bias=layer.dt_bias

621 + )

622 + initial_state = ssm_state[prefill_state_indices]

623 + initial_state[~has_initial_state, ...] = 0

624 + attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu(

625 + query=query,

626 + key=key,

627 + value=value,

628 + g=g,

629 + beta=beta,

630 + initial_state=initial_state,

631 + output_final_state=True,

632 + cu_seqlens=prefill_qsl,

633 + head_first=False,

634 + use_qk_l2norm_in_kernel=True,

635 + )

636 + ssm_state[prefill_state_indices] = last_recurrent_state.to(

637 + ssm_state.dtype, copy=False

638 + )

639 + out[:] = attn_out.squeeze(0)

640 + return out

641 +

642 +

643 def cpu_gdn_attention_core_fake(

644 mixed_qkv: torch.Tensor,

645 b: torch.Tensor,

vllm/model_executor/layers/utils.py

+7 -1

@@ -239,12 +239,18 @@ def dispatch_cpu_unquantized_gemm(

239 # For now it should be a causal_conv1d op

240 if torch.cpu._is_amx_tile_supported():

241 # prepack conv weight

242 - layer.weight.data = ops.causal_conv1d_weight_pack(

242 + unpacked = (

243 layer.weight.view(

244 layer.weight.size(0),

245 layer.weight.size(2),

246 )

247 + .contiguous()

248 + .clone()

249 )

250 + # Stash the un-packed (dim, width) weight so the speculative-decode

251 + # GDN path (which uses torch conv, not the AMX kernel) can use it.

252 + layer._cpu_unpacked_conv_weight = unpacked

253 + layer.weight.data = ops.causal_conv1d_weight_pack(unpacked)

254 return

255

256 N, K = layer.weight.size()