From 05fa8183a6e2c1e1c6e4643d54b9709fa29ff753 Mon Sep 17 00:00:00 2001 From: guybd Date: Mon, 13 Jul 2026 07:16:18 +0300 Subject: [PATCH] [CPU][Spec Decode] Support DFlash speculative decoding for GDN models on CPU (#46090) Signed-off-by: guybd Signed-off-by: Guy Boudoukh Co-authored-by: Cursor Co-authored-by: Li, Jiang --- csrc/cpu/sgl-kernels/fla.cpp | 255 ++++++++++ csrc/cpu/torch_bindings.cpp | 17 + vllm/_custom_ops.py | 34 ++ .../layers/mamba/ops/cpu/gdn_attention.py | 451 +++++++++++++++++- vllm/model_executor/layers/utils.py | 8 +- 5 files changed, 755 insertions(+), 10 deletions(-) diff --git a/csrc/cpu/sgl-kernels/fla.cpp b/csrc/cpu/sgl-kernels/fla.cpp index bf1b6444bd..b9a793c599 100644 --- a/csrc/cpu/sgl-kernels/fla.cpp +++ b/csrc/cpu/sgl-kernels/fla.cpp @@ -1116,6 +1116,164 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl( }); } +// Speculative-decode variant: processes a varlen batch where each sequence has +// ``q_len`` draft tokens, runs the recurrence sequentially over those tokens +// (inside the kernel, so one dispatch handles the whole draft block), reads the +// initial state from cache slot ``num_accepted-1`` and stores the state *after* +// 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 +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] + const scalar_t* __restrict__ v_ptr, // [T, HV, EV] + const param_t* __restrict__ A_log_ptr, + const scalar_t* __restrict__ a_ptr, // [T, HV] + const scalar_t* __restrict__ dt_bias_ptr, + const scalar_t* __restrict__ b_ptr, // [T, HV] + 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, + scalar_t* __restrict__ o_ptr, // [T, HV, EV] + float* __restrict__ qk_scale_buf, // [2, T, HK] + int64_t total_tokens, + int64_t batch_size, + int64_t spec_stride, + int64_t num_heads, + int64_t head_dim, + int64_t v_num_heads, + int64_t v_head_dim, + int64_t q_strideT, + int64_t q_strideH, + int64_t k_strideT, + int64_t k_strideH, + int64_t v_strideT, + int64_t v_strideH, + int64_t state_slot_stride, + bool use_qk_l2norm_in_kernel, + double softplus_threshold) { + 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); + + if (use_qk_l2norm_in_kernel) { + float eps = 1e-5f; + at::parallel_for(0, total_tokens * num_heads, 0, [&](int64_t begin, int64_t end) { + for (int64_t i = begin; i < end; ++i) { + int64_t ti = i / num_heads; + int64_t ni = i % num_heads; + const scalar_t* qp = q_ptr + ti * q_strideT + ni * q_strideH; + const scalar_t* kp = k_ptr + ti * k_strideT + ni * k_strideH; + float sq = 0.f, sk = 0.f; + for (int64_t d = 0; d < head_dim; ++d) { + float qv = (float)qp[d]; + sq += qv * qv; + float kv = (float)kp[d]; + sk += kv * kv; + } + qk_scale_buf[ti * num_heads + ni] = 1.f / std::sqrt(sq + eps); + qk_scale_buf[total_tokens * num_heads + ti * num_heads + ni] = 1.f / std::sqrt(sk + eps); + } + }); + } + + at::parallel_for(0, batch_size * v_num_heads, 0, [&](int64_t begin, int64_t end) { + for (int64_t idx = begin; idx < end; ++idx) { + int64_t bi = idx / v_num_heads; + int64_t ni = idx % v_num_heads; + int64_t kh = ni / group_size; + int64_t q_start = cu_seqlens_ptr[bi]; + int64_t q_len = cu_seqlens_ptr[bi + 1] - q_start; + if (q_len <= 0) { + continue; + } + int64_t acc = (int64_t)num_accepted_ptr[bi]; + // Clamp acc-1 to >=0: when num_accepted is 0 the unclamped index reads + // out of bounds and yields an arbitrary prev_slot used to index the SSM + // state. Mirrors the GPU guard tl.maximum(num_accepted - 1, 0). + int64_t prev_slot = + (int64_t)spec_indices_ptr[bi * spec_stride + (acc > 0 ? acc - 1 : 0)]; + 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; + 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); + fVec g_val_exp_vec = fVec(g_val_exp); + float beta_val = 1.f / (1.f + std::exp(-(float)b_ptr[token * v_num_heads + ni])); + fVec beta_vec = fVec(beta_val); + int64_t q_offset = token * q_strideT + kh * q_strideH; + int64_t k_offset = token * k_strideT + kh * k_strideH; + float q_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[token * num_heads + kh] : 1.f; + float k_scale = + use_qk_l2norm_in_kernel ? qk_scale_buf[total_tokens * num_heads + token * num_heads + kh] : 1.f; + int64_t v_offset = token * v_strideT + ni * v_strideH; + int64_t o_offset = (token * v_num_heads + ni) * v_head_dim; + int64_t dvi = 0; + for (; dvi <= v_head_dim - VecSize; dvi += VecSize) { + fVec kv_mem_vec0 = fVec(0.f); + 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); + 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; + } + bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi); + fVec v_vec0, v_vec1; + std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec); + fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec; + fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec; + fVec o_vec0 = fVec(0.f); + fVec o_vec1 = fVec(0.f); + 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); + 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); + } + bVec o_vec = at::vec::convert_from_float(o_vec0, o_vec1); + o_vec.store(o_ptr + o_offset + dvi); + } + for (; dvi < v_head_dim; ++dvi) { + 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; + } + float v_val = (float)v_ptr[v_offset + dvi]; + float dt_val = (v_val - kv_mem_val) * beta_val; + float o_val = 0.f; + 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; + o_val += ns * q_val * scale; + } + o_ptr[o_offset + dvi] = (scalar_t)o_val; + } + prev_slot = cur_slot; + } + } + }); +} + template void fused_gdn_gating_kernel_impl( float* __restrict__ A_log, @@ -1500,6 +1658,103 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu( 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) +// spec_state_indices: [batch, S] INT32 (S = num_spec + 1) +// num_accepted_tokens: [batch] INT32 +// cu_seqlens: [batch + 1] INT32 +// Returns output: [T, HV, EV] +at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu( + const at::Tensor& A_log, + const at::Tensor& dt_bias, + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& v, + const at::Tensor& a, + const at::Tensor& b, + at::Tensor& initial_state_source, + const at::Tensor& spec_state_indices, + const at::Tensor& num_accepted_tokens, + const at::Tensor& cu_seqlens, + bool use_qk_l2norm_in_kernel, + double softplus_beta = 1.0, + double softplus_threshold = 20.0) { + CHECK_DIM(3, q); + CHECK_DIM(3, v); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(q); + int64_t total_tokens = q.size(0); + int64_t num_heads = q.size(1); + 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_INPUT_SHAPE_DTYPE(num_accepted_tokens, {batch_size}, at::kInt); + CHECK_INPUT_SHAPE_DTYPE(cu_seqlens, {batch_size + 1}, at::kInt); + 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.size(0) >= batch_size, + "initial_state_source capacity too small: size(0)=", + initial_state_source.size(0), ", batch_size=", batch_size); + + int64_t q_strideT = q.stride(0); + int64_t q_strideH = q.stride(1); + int64_t k_strideT = k.stride(0); + int64_t k_strideH = k.stride(1); + int64_t v_strideT = v.stride(0); + int64_t v_strideH = v.stride(1); + int64_t state_slot_stride = initial_state_source.stride(0); + + 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); + }); + return o; +} + // A_log: [num_v_heads] // a: [batch, num_v_heads] // b: [batch, num_v_heads] diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index cfa296e73b..50ce23d89b 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -120,6 +120,14 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu( bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0, double softplus_threshold = 20.0); +at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu( + const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q, + const at::Tensor& k, const at::Tensor& v, const at::Tensor& a, + const at::Tensor& b, at::Tensor& initial_state_source, + const at::Tensor& spec_state_indices, const at::Tensor& num_accepted_tokens, + const at::Tensor& cu_seqlens, bool use_qk_l2norm_in_kernel, + double softplus_beta = 1.0, double softplus_threshold = 20.0); + std::tuple fused_gdn_gating_cpu( const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b, const at::Tensor& dt_bias); @@ -508,6 +516,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "softplus_threshold=20.0) -> Tensor"); ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU, &fused_sigmoid_gating_delta_rule_update_cpu); + ops.def( + "fused_sigmoid_gating_delta_rule_update_spec_cpu(Tensor A_log, Tensor " + "dt_bias, Tensor q, Tensor k, Tensor v, Tensor a, Tensor b, " + "Tensor(a!) initial_state_source, Tensor spec_state_indices, " + "Tensor num_accepted_tokens, Tensor cu_seqlens, bool " + "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float " + "softplus_threshold=20.0) -> Tensor"); + ops.impl("fused_sigmoid_gating_delta_rule_update_spec_cpu", torch::kCPU, + &fused_sigmoid_gating_delta_rule_update_spec_cpu); ops.def( "fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) " "-> (Tensor, Tensor)"); diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 2e1ce0a8de..f7c69f6200 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3252,6 +3252,40 @@ def fused_sigmoid_gating_delta_rule_update_cpu( ) +def fused_sigmoid_gating_delta_rule_update_spec_cpu( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + initial_state_source: torch.Tensor, + spec_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + cu_seqlens: torch.Tensor, + use_qk_l2norm_in_kernel: bool, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, +) -> torch.Tensor: + return torch.ops._C.fused_sigmoid_gating_delta_rule_update_spec_cpu( + A_log, + dt_bias, + q, + k, + v, + a, + b, + initial_state_source, + spec_state_indices, + num_accepted_tokens, + cu_seqlens, + use_qk_l2norm_in_kernel, + softplus_beta, + softplus_threshold, + ) + + def fused_gdn_gating_cpu( A_log: torch.Tensor, a: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index e5ef487ee9..9f3aa12c8d 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -4,6 +4,7 @@ from __future__ import annotations import torch +import torch.nn.functional as F import vllm._custom_ops as ops from vllm.forward_context import ForwardContext, get_forward_context @@ -46,11 +47,52 @@ def cpu_gdn_attention_core( if attn_metadata_i.num_actual_tokens == 0: return + assert mixed_qkv.dtype == torch.bfloat16, "CPU GDN attention requires BF16." + + # The conv-state cache temporal dim is ``width - 1`` when speculative + # decoding is disabled, and ``width - 1 + num_spec`` when it is enabled + # (vLLM allocates a wider rolling buffer to support draft-token rollback). + width = layer.conv1d.weight.size(-1) + conv_cache = layer.kv_cache[0] + if is_conv_state_dim_first(): + # DS layout: (num_slots, dim, state_len) + state_len = conv_cache.size(-1) + else: + # SD layout: (num_slots, state_len, dim) + state_len = conv_cache.size(-2) + + spec_decode_cache = state_len > (width - 1) + + if not spec_decode_cache: + # ===== Fast path: speculative decoding NOT configured ===== + # conv-state is exactly (..., dim, width-1); use the original AMX / + # torch implementation unchanged. + _cpu_gdn_attention_nonspec( + layer, attn_metadata_i, mixed_qkv, b, a, core_attn_out + ) + return + + # ===== Speculative-decoding-aware path (wide rolling conv buffer) ===== + _cpu_gdn_attention_spec_aware( + layer, attn_metadata_i, mixed_qkv, b, a, core_attn_out, width, state_len + ) + + +# --------------------------------------------------------------------------- +# Non-speculative implementation: conv-state is exactly width-1. +# --------------------------------------------------------------------------- +def _cpu_gdn_attention_nonspec( + layer, + attn_metadata_i: GDNAttentionMetadata, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + core_attn_out: torch.Tensor, +) -> None: assert ( attn_metadata_i.spec_sequence_masks is None and attn_metadata_i.num_accepted_tokens is None - ), "speculative decode not supported in CPU GDN attention." - assert mixed_qkv.dtype == torch.bfloat16, "CPU GDN attention requires BF16." + ), "speculative decode not supported without a wide conv-state cache." state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor query_start_loc = attn_metadata_i.non_spec_query_start_loc @@ -79,12 +121,7 @@ def cpu_gdn_attention_core( b = b.contiguous() num_allocated_slots, head_num, v_dim, k_dim = ssm_state.size() - ssm_state = ssm_state.view( - num_allocated_slots, - head_num, - k_dim, - v_dim, - ) + ssm_state = ssm_state.view(num_allocated_slots, head_num, k_dim, v_dim) num_decodes = attn_metadata_i.num_decodes num_decode_tokens = attn_metadata_i.num_decode_tokens @@ -109,7 +146,6 @@ def cpu_gdn_attention_core( ) else: decode_conv_state = conv_state[decode_state_indices].contiguous() - decode_mixed_qkv = causal_conv1d_update_torch( # [B, dim] -> [B, dim, 1] x=decode_mixed_qkv.unsqueeze(-1), @@ -207,6 +243,403 @@ def cpu_gdn_attention_core( core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0) +# --------------------------------------------------------------------------- +# Speculative-decoding-aware implementation (wide rolling conv buffer). +# +# Mirrors the GPU `_forward_core` spec path using torch + the existing CPU +# kernels. The conv-state is a per-sequence rolling buffer of length +# ``state_len = width - 1 + num_spec`` (single cache slot, column 0 of the +# spec block table), and the SSM recurrent state is stored multi-slot: the +# state *after* draft token ``t`` lives in block-table column ``t``, so the +# next step can resume from the slot of the last accepted token. +# --------------------------------------------------------------------------- +def _conv_buffer_view(layer) -> torch.Tensor: + """Return the conv-state cache as (num_slots, dim, state_len).""" + conv_cache = layer.kv_cache[0] + if is_conv_state_dim_first(): + return conv_cache # (num_slots, dim, state_len) + return conv_cache.transpose(-1, -2) # view -> (num_slots, dim, state_len) + + +def _ssm_state_view(layer) -> torch.Tensor: + ssm_state = layer.kv_cache[1] + num_slots, head_num, v_dim, k_dim = ssm_state.size() + return ssm_state.view(num_slots, head_num, k_dim, v_dim) + + +def _unpacked_conv_weight(layer) -> torch.Tensor: + """Return the plain (dim, width) conv weight. + + On AMX the conv1d weight is VNNI-packed in place at load time (only usable + by the AMX C++ kernel), so the torch spec-decode path relies on the + un-packed copy stashed by ``dispatch_cpu_unquantized_gemm``. + """ + w = getattr(layer.conv1d, "_cpu_unpacked_conv_weight", None) + if w is not None: + return w + w = layer.conv1d.weight + if w.dim() == 3: + return w.view(w.size(0), w.size(2)) + return w + + +def _cpu_gdn_attention_spec_aware( + layer, + attn_metadata_i: GDNAttentionMetadata, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + core_attn_out: torch.Tensor, + width: int, + state_len: int, +) -> None: + mixed_qkv = mixed_qkv.contiguous() + a = a.contiguous() + b = b.contiguous() + + spec_sequence_masks = attn_metadata_i.spec_sequence_masks + conv_buf = _conv_buffer_view(layer) # (num_slots, dim, state_len) + ssm_state = _ssm_state_view(layer) + + if spec_sequence_masks is None: + # No spec sequences in this batch (e.g. the prompt prefill step while + # speculative decoding is configured). Process as prefill/decode using + # torch conv (which only touches the first ``width-1`` columns of the + # wide buffer, leaving the rolling history untouched). + _spec_aware_nonspec( + layer, + attn_metadata_i, + mixed_qkv, + b, + a, + core_attn_out, + conv_buf, + ssm_state, + width, + ) + return + + # Split spec vs non-spec tokens. + spec_token_indx = attn_metadata_i.spec_token_indx + non_spec_token_indx = attn_metadata_i.non_spec_token_indx + num_prefills = attn_metadata_i.num_prefills + num_decodes = attn_metadata_i.num_decodes + + if num_prefills == 0 and num_decodes == 0: + mixed_qkv_spec = mixed_qkv + b_spec = b + a_spec = a + spec_out_indx = None + else: + assert spec_token_indx is not None + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx) + b_spec = b.index_select(0, spec_token_indx) + a_spec = a.index_select(0, spec_token_indx) + spec_out_indx = spec_token_indx + + spec_out = _spec_forward( + layer, + attn_metadata_i, + mixed_qkv_spec, + b_spec, + a_spec, + conv_buf, + ssm_state, + width, + state_len, + ) + + # Process the (rare) non-spec subset (prefill) tokens, if any. + nonspec_out = None + if (num_prefills > 0 or num_decodes > 0) and non_spec_token_indx is not None: + mixed_qkv_ns = mixed_qkv.index_select(0, non_spec_token_indx) + b_ns = b.index_select(0, non_spec_token_indx) + a_ns = a.index_select(0, non_spec_token_indx) + nonspec_out = _spec_aware_nonspec_subset( + layer, attn_metadata_i, mixed_qkv_ns, b_ns, a_ns, conv_buf, ssm_state, width + ) + + # Scatter outputs back to their token positions. + if spec_out_indx is None: + core_attn_out[: spec_out.size(0)] = spec_out + else: + core_attn_out.index_copy_(0, spec_out_indx, spec_out) + if nonspec_out is not None: + core_attn_out.index_copy_(0, non_spec_token_indx, nonspec_out) + + +def _spec_forward( + layer, + attn_metadata_i: GDNAttentionMetadata, + mixed_qkv_spec: torch.Tensor, + b_spec: torch.Tensor, + a_spec: torch.Tensor, + conv_buf: torch.Tensor, + ssm_state: torch.Tensor, + width: int, + state_len: int, +) -> torch.Tensor: + """Run the GDN core for the multi-query (speculative) tokens. + + Returns the core attention output for the spec tokens, in the same token + order as ``mixed_qkv_spec`` (i.e. ordered by ``spec_query_start_loc``). + """ + num_spec_decodes = attn_metadata_i.num_spec_decodes + spec_state_indices = attn_metadata_i.spec_state_indices_tensor + spec_qsl = attn_metadata_i.spec_query_start_loc + num_accepted = attn_metadata_i.num_accepted_tokens + assert spec_state_indices is not None + assert spec_qsl is not None + assert num_accepted is not None + + spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1].to("cpu", torch.int64) + num_acc_cpu = num_accepted[:num_spec_decodes].to("cpu", torch.int64) + seq_starts = spec_qsl_cpu[:-1] + seq_lens = spec_qsl_cpu[1:] - spec_qsl_cpu[:-1] + + # ---- 1. Convolution (per-sequence rolling buffer) ---- + w2d = _unpacked_conv_weight(layer) # (dim, width) + dim = w2d.size(0) + w = w2d.unsqueeze(1) # (dim, 1, width) for F.conv1d depthwise + bias = layer.conv1d.bias + silu = layer.activation == "silu" + + conv_out = torch.empty_like(mixed_qkv_spec) + col0 = spec_state_indices[:, 0].to("cpu", torch.int64) + for i in range(num_spec_decodes): + q_i = int(seq_lens[i].item()) + if q_i == 0: + continue + start = int(seq_starts[i].item()) + slot0 = int(col0[i].item()) + a_prev = int(num_acc_cpu[i].item()) + offset = a_prev - 1 + B = conv_buf[slot0] # (dim, state_len) + x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype) + prior = B[:, offset : offset + (width - 1)] + conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) # (1, dim, w-1+q) + out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i) + if silu: + out = F.silu(out) + conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype) + # Roll the buffer: drop ``a_prev`` from the front, append the new + # draft tokens, keep total length == state_len. + keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)] + new_B = torch.cat([keep, x_seq], dim=-1) + B.copy_(new_B) + + # ---- 2. Recurrent (multi-slot SSM state) ---- + # Single fused kernel call: it runs the recurrence over each sequence's + # draft tokens internally, resumes from slot ``num_accepted-1`` and stores + # the state after token ``t`` into slot ``t`` (rollback for the next step). + query, key, value = layer.rearrange_mixed_qkv(conv_out) + query = query.squeeze(0).contiguous() + key = key.squeeze(0).contiguous() + value = value.squeeze(0).contiguous() + spec_idx = spec_state_indices[:num_spec_decodes].to(torch.int32).contiguous() + num_acc = num_accepted[:num_spec_decodes].to(torch.int32).contiguous() + cu = spec_qsl[: num_spec_decodes + 1].to(torch.int32).contiguous() + out_spec = ops.fused_sigmoid_gating_delta_rule_update_spec_cpu( + A_log=layer.A_log, + dt_bias=layer.dt_bias, + q=query, + k=key, + v=value, + a=a_spec.contiguous(), + b=b_spec.contiguous(), + initial_state_source=ssm_state, + spec_state_indices=spec_idx, + num_accepted_tokens=num_acc, + cu_seqlens=cu, + use_qk_l2norm_in_kernel=True, + ) + return out_spec + + +def core_attn_out_like(layer, mixed_qkv_spec: torch.Tensor) -> torch.Tensor: + num_tokens = mixed_qkv_spec.size(0) + return torch.zeros( + (num_tokens, layer.num_v_heads // layer.tp_size, layer.head_v_dim), + dtype=mixed_qkv_spec.dtype, + device=mixed_qkv_spec.device, + ) + + +def _spec_aware_nonspec( + layer, + attn_metadata_i: GDNAttentionMetadata, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + core_attn_out: torch.Tensor, + conv_buf: torch.Tensor, + ssm_state: torch.Tensor, + width: int, +) -> None: + """Non-spec prefill/decode with a wide conv buffer (torch path).""" + state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor + query_start_loc = attn_metadata_i.non_spec_query_start_loc + assert state_indices_tensor is not None + assert query_start_loc is not None + + conv_weights = _unpacked_conv_weight(layer) + + num_decodes = attn_metadata_i.num_decodes + num_decode_tokens = attn_metadata_i.num_decode_tokens + num_prefills = attn_metadata_i.num_prefills + num_prefill_tokens = attn_metadata_i.num_prefill_tokens + + if num_decodes > 0: + decode_mixed_qkv = mixed_qkv[:num_decode_tokens] + decode_b = b[:num_decode_tokens] + decode_a = a[:num_decode_tokens] + decode_state_indices = state_indices_tensor[:num_decodes] + # Only the first ``width-1`` columns hold the real conv state. + decode_conv_state = conv_buf[decode_state_indices][ + :, :, : width - 1 + ].contiguous() + decode_mixed_qkv = causal_conv1d_update_torch( + x=decode_mixed_qkv.unsqueeze(-1), + conv_state=decode_conv_state, + weight=conv_weights, + bias=layer.conv1d.bias, + activation=layer.activation, + ).squeeze(-1) + conv_buf[decode_state_indices, :, : width - 1] = decode_conv_state + + query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) + # rearrange_mixed_qkv can return views whose last dim is not + # contiguous; the fused CPU kernel requires a contiguous last dim. + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu( + A_log=layer.A_log, + dt_bias=layer.dt_bias, + q=query, + k=key, + v=value, + a=decode_a.contiguous(), + b=decode_b.contiguous(), + initial_state_source=ssm_state, + initial_state_indices=decode_state_indices, + cu_seqlens=query_start_loc[: num_decodes + 1], + use_qk_l2norm_in_kernel=True, + ) + core_attn_out[:num_decode_tokens] = attn_out.squeeze(1) + + if num_prefills > 0: + has_initial_state = attn_metadata_i.has_initial_state + assert has_initial_state is not None + prefill_token_start = num_decode_tokens + prefill_token_end = prefill_token_start + num_prefill_tokens + prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end] + prefill_b = b[prefill_token_start:prefill_token_end] + prefill_a = a[prefill_token_start:prefill_token_end] + prefill_state_indices = state_indices_tensor[ + num_decodes : num_decodes + num_prefills + ] + prefill_query_start_loc = ( + query_start_loc[num_decodes : num_decodes + num_prefills + 1] + - num_decode_tokens + ) + prefill_has_initial_state = has_initial_state[ + num_decodes : num_decodes + num_prefills + ] + # ``causal_conv1d_torch`` only touches columns [:width-1] of the buffer. + prefill_mixed_qkv = causal_conv1d_torch( + x=prefill_mixed_qkv.transpose(0, 1), + weight=conv_weights, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_query_start_loc, + cache_indices=prefill_state_indices, + has_initial_state=prefill_has_initial_state, + activation=layer.activation, + ).transpose(0, 1) + + query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv) + g, beta = ops.fused_gdn_gating_cpu( + A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias + ) + initial_state = ssm_state[prefill_state_indices] + initial_state[~prefill_has_initial_state, ...] = 0 + attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu( + query=query, + key=key, + value=value, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=prefill_query_start_loc, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + ssm_state[prefill_state_indices] = last_recurrent_state.to( + ssm_state.dtype, copy=False + ) + core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0) + + +def _spec_aware_nonspec_subset( + layer, + attn_metadata_i: GDNAttentionMetadata, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + conv_buf: torch.Tensor, + ssm_state: torch.Tensor, + width: int, +) -> torch.Tensor: + """Process non-spec (prefill) tokens that coexist with spec sequences. + + Returns outputs ordered like ``non_spec_token_indx``. + """ + out = core_attn_out_like(layer, mixed_qkv) + has_initial_state = attn_metadata_i.has_initial_state + prefill_state_indices = attn_metadata_i.prefill_state_indices + prefill_qsl = attn_metadata_i.prefill_query_start_loc + assert prefill_state_indices is not None and prefill_qsl is not None + assert has_initial_state is not None + + conv_weights = _unpacked_conv_weight(layer) + conv_out = causal_conv1d_torch( + x=mixed_qkv.transpose(0, 1), + weight=conv_weights, + bias=layer.conv1d.bias, + conv_states=conv_buf, + query_start_loc=prefill_qsl, + cache_indices=prefill_state_indices, + has_initial_state=has_initial_state, + activation=layer.activation, + ).transpose(0, 1) + + query, key, value = layer.rearrange_mixed_qkv(conv_out) + g, beta = ops.fused_gdn_gating_cpu( + A_log=layer.A_log, a=a, b=b, dt_bias=layer.dt_bias + ) + initial_state = ssm_state[prefill_state_indices] + initial_state[~has_initial_state, ...] = 0 + attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu( + query=query, + key=key, + value=value, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=prefill_qsl, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + ssm_state[prefill_state_indices] = last_recurrent_state.to( + ssm_state.dtype, copy=False + ) + out[:] = attn_out.squeeze(0) + return out + + def cpu_gdn_attention_core_fake( mixed_qkv: torch.Tensor, b: torch.Tensor, diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 6ca42c0e7f..4d2e50420f 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -239,12 +239,18 @@ def dispatch_cpu_unquantized_gemm( # For now it should be a causal_conv1d op if torch.cpu._is_amx_tile_supported(): # prepack conv weight - layer.weight.data = ops.causal_conv1d_weight_pack( + unpacked = ( layer.weight.view( layer.weight.size(0), layer.weight.size(2), ) + .contiguous() + .clone() ) + # Stash the un-packed (dim, width) weight so the speculative-decode + # GDN path (which uses torch conv, not the AMX kernel) can use it. + layer._cpu_unpacked_conv_weight = unpacked + layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) return N, K = layer.weight.size()