From fb05c2e9a2ded4d42861bc000ac01778cd1ba4c9 Mon Sep 17 00:00:00 2001 From: markaalonzo <267525922+markaalonzo@users.noreply.github.com> Date: Tue, 28 Apr 2026 02:00:01 -0400 Subject: [PATCH] server: defer recurrent-state reset to graph build (addresses #1696 review) (#1696) ikawrakow asked the seq_rm zeroing be hoisted out of the per-cell loop and preferably done as part of the compute graph. This rewrites the fix: - llama_kv_cache gets a pending_recurrent_reset bitmap, sized to qnext_state_slots and indexed by slot index. - llama_kv_cache_seq_rm only marks the bitmap; no CPU-side tensor writes, no work inside the per-cell loop. - delta_net::build_layer_attn_linear ORs the bitmap into the existing reset_state flag, so the in-graph ggml_scale(state, 0.0f) path that already handled batch.pos[0] == 0 now also handles slot reset. - build_qwen3next and build_qwen35 clear the bitmap after ggml_build_forward_expand, once per graph build. - llama_kv_cache_clear also drops the bitmap, since the full-buffer clear it already does makes any pending in-graph reset redundant. Previous behavior is preserved (every recurrent layer column for the released slot is zeroed before the next request reuses it), but the work moves from the CPU into the compute graph, runs once per graph instead of once per cell, and reuses the existing reset path referenced in the review. --- src/graphs/build_qwen35.cpp | 16 ++++++++++++++++ src/graphs/build_qwen3next.cpp | 8 ++++++++ src/llama-context.h | 9 +++++++++ src/llama-delta-net.cpp | 13 ++++++++++--- src/llama.cpp | 18 ++++++++++++++++++ 5 files changed, 61 insertions(+), 3 deletions(-) diff --git a/src/graphs/build_qwen35.cpp b/src/graphs/build_qwen35.cpp index fb19d679..c864298c 100644 --- a/src/graphs/build_qwen35.cpp +++ b/src/graphs/build_qwen35.cpp @@ -59,6 +59,14 @@ ggml_cgraph * llm_build_context::build_qwen35moe() { ggml_build_forward_expand(gf, cur); + // Consume the recurrent-state reset flags. Any slot that was marked by + // llama_kv_cache_seq_rm has had its reset op embedded in this graph + // through delta_net's reset_state branch, so we can clear the flags now; + // the reset will fire when the graph executes. + std::fill(lctx.kv_self.pending_recurrent_reset.begin(), + lctx.kv_self.pending_recurrent_reset.end(), + false); + return gf; } @@ -139,6 +147,14 @@ ggml_cgraph * llm_build_context::build_qwen35() { ggml_build_forward_expand(gf, cur); + // Consume the recurrent-state reset flags. Any slot that was marked by + // llama_kv_cache_seq_rm has had its reset op embedded in this graph + // through delta_net's reset_state branch, so we can clear the flags now; + // the reset will fire when the graph executes. + std::fill(lctx.kv_self.pending_recurrent_reset.begin(), + lctx.kv_self.pending_recurrent_reset.end(), + false); + return gf; } diff --git a/src/graphs/build_qwen3next.cpp b/src/graphs/build_qwen3next.cpp index 3ae32419..e8f3ab91 100644 --- a/src/graphs/build_qwen3next.cpp +++ b/src/graphs/build_qwen3next.cpp @@ -85,5 +85,13 @@ ggml_cgraph * llm_build_context::build_qwen3next() { ggml_build_forward_expand(gf, cur); + // Consume the recurrent-state reset flags. Any slot that was marked by + // llama_kv_cache_seq_rm has had its reset op embedded in this graph + // through delta_net's reset_state branch, so we can clear the flags now; + // the reset will fire when the graph executes. + std::fill(lctx.kv_self.pending_recurrent_reset.begin(), + lctx.kv_self.pending_recurrent_reset.end(), + false); + return gf; } diff --git a/src/llama-context.h b/src/llama-context.h index 7b6e56cf..8b3d4473 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -62,6 +62,15 @@ struct llama_kv_cache { // When true, the delta_net graph builder will enable per-step SSM state saves bool save_per_step_ssm = false; + // Set by llama_kv_cache_seq_rm when a hybrid/recurrent slot's cell is + // fully emptied. Read once during the next graph build to inject a + // state-reset op into the recurrent layers, then cleared. Indexed by + // slot/cell index (= column in cache.s_l[layer]). Sized by + // qnext_state_slots when the cache is initialized; empty for non- + // hybrid models, in which case all reads short-circuit through the + // size() bounds check. + std::vector pending_recurrent_reset; + std::vector split_k_l; std::vector split_v_l; std::vector split_s_l; diff --git a/src/llama-delta-net.cpp b/src/llama-delta-net.cpp index 1c24fc2e..899b3a5e 100644 --- a/src/llama-delta-net.cpp +++ b/src/llama-delta-net.cpp @@ -646,8 +646,12 @@ ggml_tensor * delta_net::build_layer_attn_linear(ggml_context * ctx0, ggml_cgrap GGML_ASSERT(model.layers[il].wqkv_gate != nullptr || model.layers[il].ssm_in != nullptr); if (all_same_seq) { - bool reset_state = batch.pos != nullptr && batch.pos[0] == 0; - return build_layer_attn_linear_core(ctx0, gf, cur, lctx.inp_s_seq_qnext, inp_out_ids, token_seq_ids.front(), reset_state, il, cb); + const uint32_t state_seq_id = (uint32_t) token_seq_ids.front(); + const bool needs_recurrent_reset = + state_seq_id < lctx.kv_self.pending_recurrent_reset.size() && + lctx.kv_self.pending_recurrent_reset[state_seq_id]; + bool reset_state = (batch.pos != nullptr && batch.pos[0] == 0) || needs_recurrent_reset; + return build_layer_attn_linear_core(ctx0, gf, cur, lctx.inp_s_seq_qnext, inp_out_ids, state_seq_id, reset_state, il, cb); } GGML_ASSERT(has_unique_seq_ids && "qwen3next mixed-sequence batches require unique sequence IDs per token"); @@ -657,8 +661,11 @@ ggml_tensor * delta_net::build_layer_attn_linear(ggml_context * ctx0, ggml_cgrap ggml_tensor * cur_i = ggml_view_2d(ctx0, cur, cur->ne[0], 1, cur->nb[1], (size_t) i * cur->nb[1]); ggml_tensor * inp_s_seq_qnext_i = ggml_view_2d(ctx0, lctx.inp_s_seq_qnext, 1, 1, lctx.inp_s_seq_qnext->nb[1], (size_t) i * lctx.inp_s_seq_qnext->nb[1]); - const bool reset_state_i = batch.pos != nullptr && batch.pos[i] == 0; const uint32_t state_seq_id_i = (uint32_t) token_seq_ids[i]; + const bool needs_recurrent_reset_i = + state_seq_id_i < lctx.kv_self.pending_recurrent_reset.size() && + lctx.kv_self.pending_recurrent_reset[state_seq_id_i]; + const bool reset_state_i = (batch.pos != nullptr && batch.pos[i] == 0) || needs_recurrent_reset_i; ggml_tensor * out_i = build_layer_attn_linear_core(ctx0, gf, cur_i, inp_s_seq_qnext_i, inp_out_ids, state_seq_id_i, reset_state_i, il, cb); out = out == nullptr ? out_i : ggml_concat(ctx0, out, out_i, 1); diff --git a/src/llama.cpp b/src/llama.cpp index f7b55bbf..f5774ea3 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -851,6 +851,7 @@ static bool llama_kv_cache_init( LLAMA_LOG_WARN("%s: reducing qwen3next state slots from %u to %u to fit KV cache size\n", __func__, std::max(1, cparams.n_seq_max), qnext_state_slots); } + cache.pending_recurrent_reset.assign(qnext_state_slots, false); int n_mla = 0; const int64_t n_mtp_first_layer = n_layer - hparams.nextn_predict_layers; @@ -1597,6 +1598,14 @@ static void llama_kv_cache_clear(struct llama_kv_cache & cache) { for (auto & buf : cache.bufs) { ggml_backend_buffer_clear(buf, 0); } + + // The full clear above zeroed every recurrent state buffer, so any + // pending per-slot resets recorded by an earlier seq_rm are now + // redundant. Drop them so the next graph build does not emit a + // spurious in-graph reset op. + std::fill(cache.pending_recurrent_reset.begin(), + cache.pending_recurrent_reset.end(), + false); } static bool llama_kv_cache_seq_rm( @@ -1646,6 +1655,15 @@ static bool llama_kv_cache_seq_rm( cache.cells[i].pos = -1; if (has_qnext_state) { cache.cells[i].src = i; + // Defer the recurrent-state reset to graph build time: + // delta-net's existing reset path (ggml_scale state, 0.0f) + // does the zeroing inside the compute graph, so we just + // record which slots need reset here. The flags are + // consumed and cleared at the end of build_qwen3next / + // build_qwen35. + if ((uint32_t) i < cache.pending_recurrent_reset.size()) { + cache.pending_recurrent_reset[i] = true; + } } if (new_head == cache.size) new_head = i; }