Module decode_shape 

The identity a held decode plan is baked against.

A DecodeSession holds a graph that has already been built and optimized for one specific model. Reusing it is what makes persistent decode fast — measured at 223× on CUDA against re-planning per token — and reusing it wrongly is silent: a baked graph computing the wrong architecture returns plausible logits at full speed, with nothing to report.

Until this module existed the reuse predicate (DecodeSession::is_valid_for) keyed on (seq, max_seq_len, n_layers, cache_dtype) — pure geometry. That was sufficient while one model shape drove the path and each caller happened to hold one session per model. It is not sufficient for either of the two things now in flight:

• Different architectures, identical geometry. Qwen3 (per-head Q/K RmsNorm, sliding-window mask) and Llama at the same dims produce the same geometric key, so a plan built for one would be judged valid for the other.
• Same architecture, different weights. This one needs no new model family to be wrong: two Qwen3Models with identical config and different weights collide today, and the held graph carries the first one’s baked weight Consts.

What belongs in the key — and what must NOT

The rule is baked, not bound. A held graph bakes its op structure and its weight Consts; per-token data (token_ids, rope_cos/rope_sin, mask, the KV-write offset) is const_placeholder_like and rebound every step from the current call.

So RoPE frequencies are deliberately absent from the key, and that is not an oversight. LLaMA-3.1’s scaled frequencies reach the graph through the per-token rebind, so a session is already correct across a frequency change. The instinct is “rope differs ⇒ must be in the key”, and acting on it costs plan reuse for no correctness gain — a pure performance regression that hides behind a fully green correctness suite. Over-keying is the failure mode with no error message; that is why the tests here assert both halves — that a differing key invalidates, and that a matching one does not.

Why a counter, and not weight pointers

Model identity is a ModelInstanceId — a value from a process-global monotonic counter, minted once per constructed model. It is never recycled, so two distinct models cannot share one, ever, under any ownership arrangement.

The obvious cheaper thing is to fold the weights’ Arc addresses, and this module did that for exactly one commit. It is unsound, and the way it was unsound is the reason the counter is worth its cost:

Pointer identity is safe only while something pins the allocation. A DecodeSession holds StorageCache = HashMap<NodeId, Arc<RwLock<Storage>>> — it pins the Storage it baked, not the model’s Arc<[f32]> weight buffers, which are different allocations (and on GPU the host buffer is not retained at all). So: hold a session for model A, drop A, construct B — the allocator may hand back A’s addresses, B’s key equals A’s, and the stale plan baked with A’s weight Consts is judged valid for B.

That is reachable, not exotic: a driver serving many models with cached sessions is precisely where a session outlives its model.

The lesson generalizes past this bug. The pointer scheme was justified by a proof about lifetimes, and the proof was true of the design it was written for (folding base_cache’s own Arcs) and silently stopped being true when the input moved to model-owned weights — a change made for unrelated and good reasons. A counter that is never reused needs no proof, so it cannot be invalidated by a refactor somewhere else. Prefer that to a correct argument whose premises live in another file.

One consequence worth stating rather than discovering: a model reloaded from disk gets a fresh id and rebuilds its plan. Safe, mildly wasteful, and correct by default. Mutating a weight buffer in place under a live session does not invalidate it — decode does not do that; a training consumer must not reuse a decode session across an optimizer step.

Structs

KvAllocId

A process-unique, never-recycled identity for one KV allocation — the set of storage buffers a held decode plan is welded to.

ModelInstanceId

A process-unique, never-recycled identity for one constructed model.

ShapeKeyHasher

FNV-1a 64 over a canonical encoding of a model’s decode identity.