[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>
This commit is contained in:
guybd
co-authored by
Cursor
Li, Jiang <jiang1.li@intel.com>
Cursor
Li, Jiang <jiang1.li@intel.com>
parent
d973cce3ca
commit
05fa8183a6
@@ -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 <typename scalar_t, typename param_t>
|
||||
void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl(
|
||||
const scalar_t* __restrict__ q_ptr, // [T, HK, EK]
|
||||
const scalar_t* __restrict__ k_ptr, // [T, HK, EK]
|
||||
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<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int64_t VecSize = bVec::size();
|
||||
constexpr int64_t fVecSize = fVec::size();
|
||||
int64_t group_size = v_num_heads / num_heads;
|
||||
double scale = 1 / std::sqrt((double)head_dim);
|
||||
fVec scale_vec = fVec((float)scale);
|
||||
|
||||
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<scalar_t>(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 <typename scalar_t>
|
||||
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<true>(k, {total_tokens, num_heads, head_dim}, q.scalar_type());
|
||||
CHECK_INPUT_SHAPE_DTYPE<true>(v, {total_tokens, v_num_heads, v_head_dim}, q.scalar_type());
|
||||
CHECK_INPUT_SHAPE_DTYPE<true>(a, {total_tokens, v_num_heads}, q.scalar_type());
|
||||
CHECK_INPUT_SHAPE_DTYPE<true>(b, {total_tokens, v_num_heads}, q.scalar_type());
|
||||
CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type());
|
||||
CHECK_INPUT_SHAPE_DTYPE<true>(num_accepted_tokens, {batch_size}, at::kInt);
|
||||
CHECK_INPUT_SHAPE_DTYPE<true>(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<true>(
|
||||
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<scalar_t, param_t>(
|
||||
q.data_ptr<scalar_t>(),
|
||||
k.data_ptr<scalar_t>(),
|
||||
v.data_ptr<scalar_t>(),
|
||||
A_log.data_ptr<param_t>(),
|
||||
a.data_ptr<scalar_t>(),
|
||||
dt_bias.data_ptr<scalar_t>(),
|
||||
b.data_ptr<scalar_t>(),
|
||||
spec_state_indices.data_ptr<int32_t>(),
|
||||
num_accepted_tokens.data_ptr<int32_t>(),
|
||||
cu_seqlens.data_ptr<int32_t>(),
|
||||
initial_state_source.data_ptr<float>(),
|
||||
o.data_ptr<scalar_t>(),
|
||||
qk_scale_buf.data_ptr<float>(),
|
||||
total_tokens,
|
||||
batch_size,
|
||||
spec_stride,
|
||||
num_heads,
|
||||
head_dim,
|
||||
v_num_heads,
|
||||
v_head_dim,
|
||||
q_strideT,
|
||||
q_strideH,
|
||||
k_strideT,
|
||||
k_strideH,
|
||||
v_strideT,
|
||||
v_strideH,
|
||||
state_slot_stride,
|
||||
use_qk_l2norm_in_kernel,
|
||||
softplus_threshold);
|
||||
});
|
||||
return o;
|
||||
}
|
||||
|
||||
// A_log: [num_v_heads]
|
||||
// a: [batch, num_v_heads]
|
||||
// b: [batch, num_v_heads]
|
||||
|
||||
@@ -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<at::Tensor, at::Tensor> 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)");
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user