Fix wrong output for hybrid/recurrent models at -np > 1 (graph reuse key) (#2260)

* Fix wrong output for hybrid/recurrent models at -np > 1 (graph reuse key)

Hybrid and recurrent architectures return silently wrong output when more than
one sequence is resident. No crash and no warning: every slot keeps producing
fluent text, it is just no longer conditioned on that slot's prompt, and slots
degenerate into repetition loops. Reported in #1932.

Three things have to line up, and on these architectures they do:

  1. can_reuse_graph() keys reuse on the ubatch SHAPE. Two consecutive decode
     steps for different sequences match on every field it checks.
  2. update_cache_copies() re-points the baked view_offs for K/V, but only for
     attention layers -- it skips recurrent ones via
     is_attn_layer = !hparams.is_recurrent(il).
  3. The delta-net bakes the recurrent state row into the graph as a
     compile-time view offset, not as an input tensor.

So a graph built for sequence A is reused to decode sequence B and nothing
re-points the recurrent state: every sequence reads and writes sequence A's
state. On Qwen3.6-35B-A3B only 10 of 40 layers carry a KV cache, so the 30
layers that silently share state are three quarters of the network.

Fix: extend the reuse key with a fingerprint of the ubatch's sequence
composition -- which sequences, in what order, and which start at position 0.
That last term matters because a state reset is baked into the graph as a node;
it mirrors exactly the condition the builder itself uses (batch.pos[i] == 0 in
build_layer_attn_linear). Gated on llm_arch_is_hybrid() ||
llm_arch_is_recurrent() on both the compare and the compute side, so
architectures that never consult the fingerprint do not pay to build it.

Evidence, RTX 3090 / sm_86, Qwen3.6-35B-A3B-UD-IQ4_XS, 6 concurrent requests
each carrying a unique codeword, 600-token generations, two rounds:

  before (IK_LEGACY_GRAPH_REUSE=1):  7/12 replies degenerate
  after:                             0/12

np=1 is unaffected. Throughput is 132.9 tok/s against 132.4 before, and the
fingerprint is a pure additional invalidation -- it can only ever add reuse
misses, never remove them -- so counting the misses it causes on its own bounds
its cost exactly. Over 13000 can_reuse_graph() calls at np=1:

  calls=13000  hit=12949  miss_other=51  miss_fingerprint_only=0

Zero, so np=1 graph reuse is bit-identical to before this patch. That covers
MTP, which runs at n_parallel == 1: its draft/verify ubatch alternation was
already keyed by the existing n_tokens / mtp_op_type / mtp_step_idx /
mtp_n_heads checks, and all the fingerprint adds beyond those is per-token
seq_ids and the pos == 0 flags, both constant at np=1 during decode. prev and
prev_mtp are populated through the same reference binding, so the MTP cache
carries the fingerprint too.

IK_LEGACY_GRAPH_REUSE restores the previous behaviour so the before/after above
can be reproduced from a single build.

* Address review: fold the arch and legacy checks into the fingerprint

llama_graph_bakes_seq_state moves above the fingerprint; the fingerprint takes the arch
and returns 0 for architectures that do not bake sequence state, and when
IK_LEGACY_GRAPH_REUSE is set. Both call sites become a plain call and the two guards
live in one place.

Returning 0 keeps behaviour identical for everything else: stored and computed values
are both 0, so the comparison always matches and reuse proceeds as before this patch.

Re-verified at np=6, same binary, only the env var differing:
  legacy  6/12 sequences degenerated  -> DIRTY
  fixed   0/12                        -> clean
This commit is contained in:
ShubhamPriyadarshi 2026-08-06 19:15:20 +05:30 committed by GitHub
parent d44e2cbe57
commit 4a4a6d3c14
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 31 additions and 1 deletions

View File

@ -564,9 +564,35 @@ struct llama_context::Prev {
int32_t mtp_n_heads;
int64_t swa_w_view;
int64_t swa_win_off;
uint64_t seq_fingerprint;
ggml_cgraph * graph;
};
static inline bool llama_graph_bakes_seq_state(const llm_arch & arch) {
return llm_arch_is_hybrid(arch) || llm_arch_is_recurrent(arch);
}
static uint64_t llama_ubatch_seq_fingerprint(const llama_batch & b, const llm_arch & arch) {
static const bool legacy_graph_reuse = getenv("IK_LEGACY_GRAPH_REUSE") != nullptr;
if (!llama_graph_bakes_seq_state(arch) || legacy_graph_reuse) {
return 0ull;
}
uint64_t h = 1469598103934665603ull;
auto mix = [&h](uint64_t v) { h ^= v; h *= 1099511628211ull; };
mix((uint64_t) (uint32_t) b.n_tokens);
for (int32_t i = 0; i < b.n_tokens; ++i) {
const int32_t ns = b.n_seq_id ? b.n_seq_id[i] : 0;
mix((uint64_t) (uint32_t) ns);
if (b.seq_id && b.seq_id[i]) {
for (int32_t j = 0; j < ns; ++j) {
mix((uint64_t) (uint32_t) b.seq_id[i][j]);
}
}
mix(b.pos && b.pos[i] == 0 ? 1ull : 0ull);
}
return h;
}
void llama_context::reset_scheduler() {
ggml_backend_sched_reset(sched);
prev.reset();
@ -597,6 +623,9 @@ bool llama_context::can_reuse_graph(const llama_batch & u_batch) {
auto the_prev = cparams.mtp_op_type == MTP_OP_NONE ? prev.get() : prev_mtp.get();
if (!the_prev || !the_prev->graph) return false;
if (u_batch.embd) return false;
if (llama_ubatch_seq_fingerprint(u_batch, model.arch) != the_prev->seq_fingerprint) {
return false;
}
auto & kv_self_used = (model.arch == LLM_ARCH_GEMMA4_MTP || model.arch == LLM_ARCH_GEMMA4_ASSISTANT) &&
mtp_target_ctx != nullptr ? mtp_target_ctx->kv_self : kv_self;
if (the_prev->save_per_step_ssm != kv_self_used.save_per_step_ssm ||
@ -6323,7 +6352,8 @@ static int llama_decode_internal(
kv_self_used.save_per_step_ssm, kv_self_used.ckpt.per_step_max_allocated,
cparams.mtp_op_type, lctx.mtp_step_idx, lctx.mtp_n_heads,
lctx.swa_window_view.w_view,
lctx.swa_window_view.win_off, gf});
lctx.swa_window_view.win_off,
llama_ubatch_seq_fingerprint(u_batch, model.arch), gf});
}
} else {
//printf("Reusing graph with type = %d, n_kv = %d, n_tokens = %d\n", cparams.mtp_op_type, (int)prev->n_kv, (int)prev->n_tokens);