From b009b797278b2a89cfe3eea024c17d16e8f81da1 Mon Sep 17 00:00:00 2001 From: Simon Teixidor Date: Wed, 9 Sep 2026 22:24:48 +0200 Subject: [PATCH 1/2] llama: flatten the GDN output projection input to 2-D build_layer_attn_linear reshapes the gated-norm result to [d, n_seq_tokens, n_seqs] before the ssm_out projection and then flattens the projection result to [n_embd, n_seq_tokens*n_seqs] on the very next line. The 3-D shape therefore carries no information past the mul_mat, but it does decide how the mul_mat is dispatched: src1 with ne12 = n_seqs is a batch of n_seqs matmuls of ne11 = n_seq_tokens columns each, and every one of them re-reads the whole ssm_out weight. Reshaping to [d, n_seq_tokens*n_seqs] instead gives a single matmul of ne11 = n_seq_tokens*n_seqs columns over one weight read. Mathematically identical: mul_mat is per-column, attn_out_norm is contiguous (the 3-D reshape already required that), and the columns are in the same order in both layouts, so the flattened call computes the same columns in the same order. The reshape_2d after the projection becomes shape-preserving. At n_seqs = 6 this is the difference between a batched call and six unbatched ones on every backend, over 48 of these projections per decode step on the 27B models. Note this only helps n_seqs > 1: one sequence with n_seq_tokens speculative tokens already arrives as a single ne11 = n_seq_tokens matmul. Only qwen35 and qwen35moe are changed, the architectures this was tested on. Assisted-by: Claude Opus 5 --- src/models/qwen35.cpp | 4 ++-- src/models/qwen35moe.cpp | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 746719da2edf..73eba2cf6c8f 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -466,8 +466,8 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( // Apply gated normalization: self.norm(core_attn_out, z) ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il); - // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim] - ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs); + // 2-D: with a [d, n_seq_tokens, n_seqs] src1 the projection is n_seqs matmuls that each read the whole ssm_out weight + ggml_tensor * final_output = ggml_reshape_2d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens * n_seqs); cb(final_output, "final_output", il); // Output projection diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index b6b5283303c6..cc7f44e56980 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -483,8 +483,8 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn_linear( // Apply gated normalization: self.norm(core_attn_out, z) ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il); - // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim] - ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs); + // 2-D: with a [d, n_seq_tokens, n_seqs] src1 the projection is n_seqs matmuls that each read the whole ssm_out weight + ggml_tensor * final_output = ggml_reshape_2d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens * n_seqs); cb(final_output, "final_output", il); // Output projection From e3f31cf6d1e7613d5c10caee3477866b6ecc21eb Mon Sep 17 00:00:00 2001 From: Simon Teixidor Date: Wed, 2 Sep 2026 20:27:19 +0200 Subject: [PATCH 2/2] qwen35, qwen35moe: run gated_delta_net in place on the recurrent state cache build_rs gains an allow_inplace option: when the sequence-copy mapping is the identity, the state input is a view of the cache slot rather than a get_rows copy. qwen35 and qwen35moe both opt in, since they build the same Gated DeltaNet graph through llm_build_delta_net_base. s_copy_main_identity joins the graph-reuse key: a seq_cp between steps changes whether the mapping is the identity, so a graph built for one must not be reused for the other. test-recurrent-state-rollback gains a check that decodes two sequences in shared ubatches while seq_cp swaps them, forks one and rolls it back, with graph reuse on. The logits are compared with a context without graph reuse (bit-exact), one with n_rs_seq = 1 and one that decodes one sequence per ubatch. It uses fresh random weights at scale 1: at the 0.01 scale of the generated models the recurrent branch does not change the logits, so a wrong state would not be seen. The check runs for qwen35 and qwen35moe, and the test now also runs on the generated qwen35moe model. Assisted-by: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01MkuKwgPNgpAg7EbxKR2iWC --- src/llama-graph.cpp | 29 ++- src/llama-graph.h | 11 +- src/llama-memory-recurrent.cpp | 14 ++ src/llama-memory-recurrent.h | 3 + src/models/qwen35.cpp | 2 +- src/models/qwen35moe.cpp | 2 +- tests/CMakeLists.txt | 9 + tests/test-recurrent-state-rollback.cpp | 260 ++++++++++++++++++++++++ 8 files changed, 321 insertions(+), 9 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 010c40bcc500..e5607ae0501f 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -358,6 +358,8 @@ bool llm_graph_input_rs::can_reuse(const llm_graph_params & params) { res &= head == mctx->get_head(); res &= rs_z == mctx->get_rs_z(); + res &= s_copy_main_identity == mctx->is_s_copy_main_identity(params.ubatch.n_seqs); + return res; } @@ -1134,6 +1136,8 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) { res &= inp_rs->head == mctx->get_recr()->get_head(); res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z(); + res &= inp_rs->s_copy_main_identity == mctx->get_recr()->is_s_copy_main_identity(params.ubatch.n_seqs); + return res; } @@ -1177,6 +1181,8 @@ bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) { res &= inp_rs->head == mctx->get_recr()->get_head(); res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z(); + res &= inp_rs->s_copy_main_identity == mctx->get_recr()->is_s_copy_main_identity(params.ubatch.n_seqs); + return res; } @@ -1265,6 +1271,8 @@ bool llm_graph_input_mem_hybrid_iswa::can_reuse(const llm_graph_params & params) res &= inp_rs->head == mctx->get_recr()->get_head(); res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z(); + res &= inp_rs->s_copy_main_identity == mctx->get_recr()->is_s_copy_main_identity(params.ubatch.n_seqs); + return res; } @@ -3489,7 +3497,8 @@ ggml_tensor * llm_graph_context::build_rs( uint32_t rs_head, uint32_t rs_size, int32_t rs_zero, - const llm_graph_get_rows_fn & get_state_rows) const { + const llm_graph_get_rows_fn & get_state_rows, + bool main_inplace) const { GGML_UNUSED(rs_size); ggml_tensor * states = ggml_reshape_2d(ctx0, s, state_size, s->ne[1]); @@ -3502,8 +3511,14 @@ ggml_tensor * llm_graph_context::build_rs( // copy states // NOTE: assuming the copy destinations are ALL contained between rs_head and rs_head + n_rs // {state_size, rs_size} -> {state_size, n_seqs} - ggml_tensor * output_states = get_state_rows(ctx0, states, state_copy_main); - ggml_build_forward_expand(gf, output_states); + ggml_tensor * output_states; + if (main_inplace) { + // rows are already in place; not expanded here, so the view stays next to its consumer + output_states = ggml_view_2d(ctx0, states, state_size, n_seqs, states->nb[1], rs_head*states->nb[1]); + } else { + output_states = get_state_rows(ctx0, states, state_copy_main); + ggml_build_forward_expand(gf, output_states); + } // copy extra states which won't be changed further (between n_seqs and n_rs) ggml_tensor * states_extra = ggml_get_rows(ctx0, states, state_copy_extra); @@ -3534,6 +3549,7 @@ static std::unique_ptr build_rs_inp_impl( inp->head = mctx_cur->get_head(); inp->rs_z = mctx_cur->get_rs_z(); + inp->s_copy_main_identity = mctx_cur->is_s_copy_main_identity(n_seqs); return inp; } @@ -3551,12 +3567,15 @@ ggml_tensor * llm_graph_context::build_rs( ggml_tensor * s, int32_t state_size, int32_t n_seqs, - const llm_graph_get_rows_fn & get_state_rows) const { + const llm_graph_get_rows_fn & get_state_rows, + bool allow_inplace) const { const auto * kv_state = inp->mctx; + const bool main_inplace = allow_inplace && inp->s_copy_main_identity; + return build_rs(s, inp->s_copy_main, inp->s_copy_extra, state_size, n_seqs, kv_state->get_n_rs(), kv_state->get_head(), kv_state->get_size(), kv_state->get_rs_z(), - get_state_rows); + get_state_rows, main_inplace); } ggml_tensor * llm_graph_context::build_rwkv_token_shift_load( diff --git a/src/llama-graph.h b/src/llama-graph.h index ce68832ce47a..35928daab07c 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -280,6 +280,9 @@ class llm_graph_input_rs : public llm_graph_input_i { // used in view offsets, need to match for valid graph reuse uint32_t head; int32_t rs_z; + + // part of the reuse key: a seq_cp between two ubatches can change it + bool s_copy_main_identity = false; }; class llm_graph_input_cross_embd : public llm_graph_input_i { @@ -1332,16 +1335,20 @@ struct llm_graph_context { uint32_t rs_head, uint32_t rs_size, int32_t rs_zero, - const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const; + const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows, + bool main_inplace = false) const; llm_graph_input_rs * build_rs_inp() const; + // allow_inplace: when no state is copied, return a view of the cache rows instead of a get_rows copy + // the consumer must read the whole state before the write-back cpy ggml_tensor * build_rs( llm_graph_input_rs * inp, ggml_tensor * s, int32_t state_size, int32_t n_seqs, - const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const; + const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows, + bool allow_inplace = false) const; ggml_tensor * build_rwkv_token_shift_load( llm_graph_input_rs * inp, diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 57919accf095..93efed585990 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -1303,6 +1303,20 @@ ggml_tensor * llama_memory_recurrent_context::get_p_l(int32_t il) const { return mem->p_l[il]; } +bool llama_memory_recurrent_context::is_s_copy_main_identity(uint32_t n_seqs) const { + // with n_rs_seq == 0, s_copy(i) is cells[head + i].src0 + if (is_full || mem->n_rs_seq > 0 || n_seqs > get_n_rs()) { + return false; + } + for (uint32_t i = 0; i < n_seqs; ++i) { + const uint32_t cell = mem->head + i; + if (mem->cells[cell].src0 != (int32_t) cell) { + return false; + } + } + return true; +} + int32_t llama_memory_recurrent_context::s_copy(int i) const { const uint32_t cell_idx = i + mem->head; const int32_t src0 = mem->cells[cell_idx].src0; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 4abb3f5cf5c0..3391403188f9 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -177,6 +177,9 @@ class llama_memory_recurrent_context : public llama_memory_context_i { int32_t s_copy(int i) const; + // s_copy(i) == head + i for all i < n_seqs, without side effects; always false with n_rs_seq > 0 + bool is_s_copy_main_identity(uint32_t n_seqs) const; + private: const llama_memory_status status; diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 73eba2cf6c8f..1b0b47f07c00 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -397,7 +397,7 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il); - ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs, ggml_get_rows, /*allow_inplace=*/true); state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); cb(state, "state_predelta", il); diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index cc7f44e56980..74ee8c0d4ac6 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -414,7 +414,7 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn_linear( ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il); - ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs, ggml_get_rows, /*allow_inplace=*/true); state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); cb(state, "state_predelta", il); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2e1f67df181c..ee05598a6aad 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -247,6 +247,15 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) set_tests_properties(test-recurrent-state-rollback-kimi-k3 PROPERTIES FIXTURES_REQUIRED generate-models ) + llama_test( + test-recurrent-state-rollback + NAME test-recurrent-state-rollback-qwen35moe + LABEL main + ARGS -m "${MODEL_DIR}/qwen35moe-moe.gguf" + ) + set_tests_properties(test-recurrent-state-rollback-qwen35moe PROPERTIES + FIXTURES_REQUIRED generate-models + ) # Test state save/load functionality across all architectures, using the generated dummy models llama_test( diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index ef05de67d004..397edc3b7b08 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -1,6 +1,7 @@ #include "arg.h" #include "common.h" #include "ggml-backend.h" +#include "gguf.h" #include "llama.h" #include "../src/llama-io.h" @@ -10,8 +11,13 @@ #include #include #include +#include +#include +#include #include +#include #include +#include #include static bool decode_tokens(llama_context * ctx, const std::vector & tokens, uint32_t count) { @@ -262,6 +268,241 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode return true; } +static void set_tensor_data_scaled(ggml_tensor * tensor, void * userdata) { + const size_t seed = *(const size_t *) userdata ^ std::hash{}(tensor->name); + std::mt19937 gen(seed); + std::normal_distribution dis(0.0f, 1.0f); + + // norm weights and matrices at scale 1: at the 0.01 scale of the generated models the recurrent + // branch is too small to change the logits, so a wrong recurrent state would go unnoticed + const bool is_norm = strstr(tensor->name, "norm") != nullptr; + const float scale = ggml_n_dims(tensor) > 1 ? 1.0f : 1.0e-2f; + + GGML_ASSERT(tensor->type == GGML_TYPE_F32); + std::vector tmp(ggml_nelements(tensor)); + for (auto & x : tmp) { + x = is_norm ? 1.0f : scale*dis(gen); + } + ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor)); +} + +// same hparams as the model file, new random weights +static llama_model * load_model_scaled(common_params & params) { + gguf_init_params gparams = { + /*.no_alloc =*/ true, + /*.ctx =*/ nullptr, + }; + gguf_context * meta_file = gguf_init_from_file(params.model.path.c_str(), gparams); + if (meta_file == nullptr) { + return nullptr; + } + gguf_context * meta = gguf_init_empty(); + gguf_set_kv(meta, meta_file); + gguf_free(meta_file); + + size_t seed = 1234; + llama_model * model = llama_model_init_from_user(meta, set_tensor_data_scaled, &seed, common_model_params_to_llama(params)); + gguf_free(meta); + return model; +} + +static void set_graph_reuse_disable(bool disable) { +#ifdef _WIN32 + _putenv_s("LLAMA_GRAPH_REUSE_DISABLE", disable ? "1" : ""); +#else + if (disable) { + setenv("LLAMA_GRAPH_REUSE_DISABLE", "1", 1); + } else { + unsetenv("LLAMA_GRAPH_REUSE_DISABLE"); + } +#endif +} + +// Two sequences share ubatches while seq_cp swaps them, forks one and rolls it back, with graph reuse on. +// With n_rs_seq = 0 the recurrent state can be read in place when no state is copied, so a seq_cp between +// two same-shape ubatches must stop the graph from being reused. Checked against a context without graph +// reuse (same ubatches, so bit-exact), a context with n_rs_seq = 1, which never reads the state in place, +// and a context that decodes one sequence per ubatch. +static bool test_seq_cp_graph_reuse(const common_params & params, llama_model * model) { + const char * func = __func__; + const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + + const auto make = [&](uint32_t n_rs_seq, bool reuse) { + auto cparams = common_context_params_to_llama(params); + cparams.n_seq_max = 3; + cparams.n_rs_seq = n_rs_seq; + cparams.n_ctx = 256; + cparams.n_batch = 64; + cparams.n_ubatch = 64; + cparams.kv_unified = true; + set_graph_reuse_disable(!reuse); + llama_context * ctx = llama_init_from_model(model, cparams); + set_graph_reuse_disable(false); + return ctx; + }; + + struct arm { + const char * name; + llama_context * ctx; + bool split; // one llama_decode per sequence + float eps; // max |logit diff| against the first arm, relative to its max |logit| + float diff_max; + }; + std::vector arms = { + { "reuse", make(0, true), false, 0.0f, 0.0f }, + { "no-reuse", make(0, false), false, 1e-6f, 0.0f }, + { "no-in-place", make(1, false), false, 1e-5f, 0.0f }, + { "seq-split", make(0, false), true, 1e-4f, 0.0f }, + }; + + const auto cleanup = [&]() { + for (auto & a : arms) { + llama_free(a.ctx); + } + }; + + for (const auto & a : arms) { + if (a.ctx == nullptr) { + fprintf(stderr, "%s : failed to init %s context\n", __func__, a.name); + cleanup(); + return false; + } + } + + llama_pos pos[3] = { 0, 0, 0 }; + uint32_t n_step = 0; + bool ok = true; + + const auto decode = [&](const std::vector & seqs, uint32_t n_tokens) { + if (!ok) { + return; + } + n_step++; + const auto token = [&](llama_seq_id s, uint32_t i) { + return (llama_token) ((13*n_step + 7*(uint32_t) s + 3*i + 1) % (uint32_t) n_vocab); + }; + + const size_t n_rows = seqs.size()*n_tokens; + std::vector> logits(arms.size()); + llama_batch batch = llama_batch_init(n_rows, 0, 1); + for (size_t k = 0; ok && k < arms.size(); ++k) { + logits[k].reserve(n_rows*n_vocab); + const size_t n_batches = arms[k].split ? seqs.size() : 1; + for (size_t b = 0; ok && b < n_batches; ++b) { + common_batch_clear(batch); + for (size_t j = 0; j < seqs.size(); ++j) { + if (arms[k].split && j != b) { + continue; + } + for (uint32_t i = 0; i < n_tokens; ++i) { + common_batch_add(batch, token(seqs[j], i), pos[seqs[j]] + (llama_pos) i, { seqs[j] }, true); + } + } + ok = llama_decode(arms[k].ctx, batch) == 0; + for (int32_t i = 0; ok && i < batch.n_tokens; ++i) { + const float * l = llama_get_logits_ith(arms[k].ctx, i); + logits[k].insert(logits[k].end(), l, l + n_vocab); + } + } + if (!ok) { + fprintf(stderr, "%s : %s decode failed at step %u\n", func, arms[k].name, n_step); + } + } + llama_batch_free(batch); + + for (size_t r = 0; ok && r < n_rows; ++r) { + const float * ref = logits[0].data() + r*n_vocab; + float ref_max = 0.0f; + for (int t = 0; t < n_vocab; ++t) { + ref_max = std::max(ref_max, std::fabs(ref[t])); + } + for (size_t k = 1; k < arms.size(); ++k) { + const float * cur = logits[k].data() + r*n_vocab; + float diff = 0.0f; + for (int t = 0; t < n_vocab; ++t) { + diff = std::max(diff, logit_diff(ref[t], cur[t])); + } + diff /= std::max(ref_max, std::numeric_limits::min()); + arms[k].diff_max = std::max(arms[k].diff_max, diff); + if (!(diff <= arms[k].eps)) { + fprintf(stderr, "%s : step %u, seq %d row %zu: %s differs from %s by %g (rel)\n", + func, n_step, seqs[r/n_tokens], r % n_tokens, arms[k].name, arms[0].name, (double) diff); + ok = false; + } + } + } + for (llama_seq_id s : seqs) { + pos[s] += (llama_pos) n_tokens; + } + }; + + const auto seq_cp = [&](llama_seq_id src, llama_seq_id dst) { + for (auto & a : arms) { + llama_memory_t mem = llama_get_memory(a.ctx); + ok = ok && llama_memory_seq_rm(mem, dst, -1, -1); + llama_memory_seq_cp(mem, src, dst, -1, -1); + } + pos[dst] = pos[src]; + }; + + const auto seq_rm = [&](llama_seq_id s) { + for (auto & a : arms) { + ok = ok && llama_memory_seq_rm(llama_get_memory(a.ctx), s, -1, -1); + } + pos[s] = 0; + }; + + const auto swap_01 = [&]() { + seq_cp(0, 2); + seq_cp(1, 0); + seq_cp(2, 1); + seq_rm(2); + }; + + // seq 1 is decoded first, so the first shared ubatch reorders the cells + decode({ 1 }, 5); + decode({ 0 }, 5); + for (uint32_t n_tokens : { 2u, 3u, 1u }) { + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + + // same ubatch shape, head and rs_z as the step before: only the state copy map changes + swap_01(); + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + + decode({ 0 }, n_tokens); + decode({ 1 }, n_tokens); + decode({ 0 }, n_tokens); + decode({ 1 }, n_tokens); + + // fork seq 0, advance it, then roll it back to the fork + seq_cp(0, 2); + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + seq_cp(2, 0); + decode({ 0, 1 }, n_tokens); + seq_rm(2); + decode({ 0, 1 }, n_tokens); + decode({ 0, 1 }, n_tokens); + } + + const int32_t n_reused = llama_perf_context(arms[0].ctx).n_reused; + if (ok && n_reused == 0) { + fprintf(stderr, "%s : graph reuse was not exercised\n", __func__); + ok = false; + } + + if (ok) { + fprintf(stderr, "%s : %u steps matched, %d graphs reused (max rel diff: %s %g, %s %g, %s %g)\n", __func__, n_step, n_reused, + arms[1].name, (double) arms[1].diff_max, arms[2].name, (double) arms[2].diff_max, arms[3].name, (double) arms[3].diff_max); + } + + cleanup(); + return ok; +} + static int test_rollback(const common_params & params, llama_model * model, uint8_t fill) { const llama_vocab * vocab = llama_model_get_vocab(model); const int n_vocab = llama_vocab_n_tokens(vocab); @@ -463,5 +704,24 @@ int main(int argc, char ** argv) { } } + // the in-place recurrent state is only enabled for qwen35 and qwen35moe + char arch[64] = {}; + llama_model_meta_val_str(model, "general.architecture", arch, sizeof(arch)); + if (strcmp(arch, "qwen35") != 0 && strcmp(arch, "qwen35moe") != 0) { + fprintf(stderr, "%s : skipping test_seq_cp_graph_reuse for %s\n", __func__, arch); + return 0; + } + + llama_model * model_scaled = load_model_scaled(params); + if (model_scaled == nullptr) { + fprintf(stderr, "%s : failed to create scaled model\n", __func__); + return 1; + } + const bool ok = test_seq_cp_graph_reuse(params, model_scaled); + llama_model_free(model_scaled); + if (!ok) { + return 1; + } + return 0; }