From 1beaaa002d22a6849653351a47f9ef9415d1d010 Mon Sep 17 00:00:00 2001 From: dmaivel Date: Sat, 2 May 2026 22:14:56 -0700 Subject: [PATCH] speculative: enable MTP per-step checkpoints with CPU recurrent layers (#1724) --- ggml/src/ggml.c | 14 ++++++++---- ggml/src/iqk/iqk_mul_mat.cpp | 41 +++++++++++++++++++++++------------- ggml/src/iqk/iqk_mul_mat.h | 2 +- src/llama.cpp | 9 ++++---- 4 files changed, 42 insertions(+), 24 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 4454ffb9..b7db11d9 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -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; + } } } diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp index 5a3800e8..b42a51ba 100644 --- a/ggml/src/iqk/iqk_mul_mat.cpp +++ b/ggml/src/iqk/iqk_mul_mat.cpp @@ -1438,7 +1438,7 @@ template 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 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(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; } diff --git a/ggml/src/iqk/iqk_mul_mat.h b/ggml/src/iqk/iqk_mul_mat.h index ae35df8e..298b00ac 100644 --- a/ggml/src/iqk/iqk_mul_mat.h +++ b/ggml/src/iqk/iqk_mul_mat.h @@ -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 } diff --git a/src/llama.cpp b/src/llama.cpp index 12299766..9efc87d7 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -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;