nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

speculative: enable MTP per-step checkpoints with CPU recurrent layers (#1724)

This commit is contained in:
dmaivel 2026-05-02 22:14:56 -07:00 committed by GitHub
parent b8eb8ccbb5
commit 1beaaa002d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 42 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
ggml/src/ggml.c
View File

@ -22662,18 +22662,17 @@ static void ggml_compute_forward_delta_net_f32(
const int nth = params->nth;
int repeat_type = dst->op_params[0];
// save_all_steps is handled by the CUDA backend only;
// the CPU path always writes to the single state slot after the output.
const int save_all_steps = dst->op_params[1];
const int64_t state_step_stride = head_dim * head_dim * n_heads * n_seqs;
float * state_working = out_data + output_size;
if (iqk_fused_delta_net(head_dim, n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs,
src2->nb[1]/sizeof(float), src2->nb[2]/sizeof(float), src2->nb[3]/sizeof(float),
q_data, k_data, v_data, g_data, beta_data, state_in,
out_data, state_working, ith, nth)) {
out_data, state_working, save_all_steps, (int) state_step_stride, ith, nth)) {
return;
}
// TODO: fix this in case we need to fall back to it.
const int64_t total_heads = n_heads * n_seqs;
const int64_t heads_per_thread = (total_heads + nth - 1) / nth;
const int64_t h_start = ith * heads_per_thread;
@ -22703,6 +22702,7 @@ static void ggml_compute_forward_delta_net_f32(
}
float * state = state_working + state_head_offset;
const int64_t state_head_size = head_dim * head_dim;
for (int64_t t = 0; t < n_tokens; ++t) {
const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride;
@ -22757,6 +22757,12 @@ static void ggml_compute_forward_delta_net_f32(
state[row + col * head_dim] = fminf(fmaxf(s, -1e6f), 1e6f);
}
}
if (save_all_steps && t + 1 < n_tokens) {
float * next_state = state_working + (t + 1) * state_step_stride + state_head_offset;
memcpy(next_state, state, state_head_size * sizeof(float));
state = next_state;
}
}
}

41
ggml/src/iqk/iqk_mul_mat.cpp
View File

@ -1438,7 +1438,7 @@ template <int head_dim>
void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs,
size_t vnb1, size_t vnb2, size_t vnb3,
const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
const float * state_in, float * out_data, float * state_out, int ith, int nth) {
const float * state_in, float * out_data, float * state_out, int save_all_steps, int state_step_stride, int ith, int nth) {
const int total_heads = n_heads * n_seqs;
const int heads_per_thread = (total_heads + nth - 1) / nth;
const int h_start = ith * heads_per_thread;
@ -1465,11 +1465,10 @@ void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type,
const int out_head_offset = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim;
const int out_token_stride = head_dim * n_heads;
for (int i = 0; i < head_dim * head_dim; ++i) {
state_out[state_head_offset + i] = state_in[state_head_offset + i];
}
float * state = state_out + state_head_offset;
for (int i = 0; i < head_dim * head_dim; ++i) {
state[i] = state_in[state_head_offset + i];
}
for (int t = 0; t < n_tokens; ++t) {
const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride;
@ -1537,6 +1536,12 @@ void iqk_fused_delta_net_neon_impl(int n_heads, int gqa_ratio, int repeat_type,
}
}
}
if (save_all_steps && t + 1 < n_tokens) {
float * next_state = state_out + (t + 1) * state_step_stride + state_head_offset;
std::memcpy(next_state, state, head_dim * head_dim * sizeof(float));
state = next_state;
}
}
}
}
@ -1545,10 +1550,10 @@ template <int head_dim>
void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs,
size_t vnb1, size_t vnb2, size_t vnb3,
const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
const float * state_in, float * out_data, float * state_out, int ith, int nth) {
const float * state_in, float * out_data, float * state_out, int save_all_steps, int state_step_stride, int ith, int nth) {
#ifdef __ARM_NEON
iqk_fused_delta_net_neon_impl<head_dim>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3,
q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, ith, nth);
q_data, k_data, v_data, g_data, beta_data, state_in, out_data, state_out, save_all_steps, state_step_stride, ith, nth);
return;
#endif
const int total_heads = n_heads * n_seqs;
@ -1581,11 +1586,10 @@ void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n
const int out_head_offset = batch_idx * (head_dim * n_heads * n_tokens) + head_idx * head_dim;
const int out_token_stride = head_dim * n_heads;
for (int i = 0; i < head_dim * head_dim; ++i) {
state_out[state_head_offset + i] = state_in[state_head_offset + i];
}
float * state = state_out + state_head_offset;
for (int i = 0; i < head_dim * head_dim; ++i) {
state[i] = state_in[state_head_offset + i];
}
for (int t = 0; t < n_tokens; ++t) {
const float * q_t = q_data + qkv_head_offset_kq + t * qkv_token_stride;
@ -1706,6 +1710,12 @@ void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n
}
#endif
#endif
if (save_all_steps && t + 1 < n_tokens) {
float * next_state = state_out + (t + 1) * state_step_stride + state_head_offset;
std::memcpy(next_state, state, head_dim * head_dim * sizeof(float));
state = next_state;
}
}
}
}
@ -1714,16 +1724,16 @@ void iqk_fused_delta_net_impl(int n_heads, int gqa_ratio, int repeat_type, int n
bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs,
size_t vnb1, size_t vnb2, size_t vnb3,
const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
const float * state_in, float * out_data, float * state_out, int ith, int nth) {
const float * state_in, float * out_data, float * state_out, int save_all_steps, int state_step_stride, int ith, int nth) {
if (head_dim != 64 && head_dim != 128) {
return false;
}
if (head_dim == 64) {
iqk_fused_delta_net_impl<64>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, q_data, k_data, v_data, g_data, beta_data, state_in,
out_data, state_out, ith, nth);
out_data, state_out, save_all_steps, state_step_stride, ith, nth);
} else {
iqk_fused_delta_net_impl<128>(n_heads, gqa_ratio, repeat_type, n_tokens, n_seqs, vnb1, vnb2, vnb3, q_data, k_data, v_data, g_data, beta_data, state_in,
out_data, state_out, ith, nth);
out_data, state_out, save_all_steps, state_step_stride, ith, nth);
}
return true;
}
@ -1762,8 +1772,9 @@ extern "C" IQK_API bool iqk_moe_fused_up_gate(long /*Nx*/, long /*Ny*/, long /*n
}
bool iqk_fused_delta_net(int, int, int, int, int, int,
size_t, size_t, size_t,
const float *, const float *, const float *, const float *, const float *,
const float *, float *, float *, int, int) {
const float *, float *, float *, int, int, int, int) {
return false;
}

2
ggml/src/iqk/iqk_mul_mat.h
View File

@ -76,7 +76,7 @@ IQK_API void iqk_topk_moe(int n_experts, int n_experts_used, int nrows, const fl
IQK_API bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs,
size_t vnb1, size_t vnb2, size_t vnb3,
const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data,
const float * state_in, float * out_data, float * state_out, int ith, int nth);
const float * state_in, float * out_data, float * state_out, int save_all_steps, int state_step_stride, int ith, int nth);
#ifdef __cplusplus
}

9
src/llama.cpp
View File

@ -6863,7 +6863,8 @@ void llama_kv_cache_clear(struct llama_context * ctx) {
// Unified speculative-checkpoint
static bool spec_ckpt_try_per_step(llama_kv_cache & kv, const llama_model & model, int max_tokens) {
// Graph-split tensors and mixed CPU/GPU configurations are not supported.
// Graph-split recurrent tensors are not supported. Mixed CPU/GPU recurrent
// placement is allowed as long as each layer has a concrete backend buffer.
bool has_gpu = false;
bool has_cpu = false;
for (const auto * sl : kv.s_l) {
@ -6878,9 +6879,9 @@ static bool spec_ckpt_try_per_step(llama_kv_cache & kv, const llama_model & mode
has_cpu = true;
}
}
if (!has_gpu || has_cpu) {
if (has_cpu && has_gpu) {
LLAMA_LOG_INFO("%s: per-step disabled — mixed CPU/GPU recurrent layers\n", __func__);
if (!has_gpu) {
if (has_cpu) {
LLAMA_LOG_INFO("%s: per-step disabled — recurrent layers are CPU-only\n", __func__);
}
kv.save_per_step_ssm = false;
return false;