Add StaticCache: fixed-capacity in-place KV cache for WebGPU

[kylo5aby]

Adds an opt-in StaticCache that eliminates per-step KV cache reallocation on WebGPU.

Fixes #1741

With DynamicCache (the default), every decode step replaces every KV tensor with a slightly larger one. On WebGPU each replacement destroys and re-creates a zero-filled GPU buffer, so the per-step cost grows with context length and jumps permanently at the buffer-pool bucket boundaries, producing a staircase-shaped latency curve.

StaticCache allocates each cache entry once at a fixed max_cache_len, and binds the model's present.* outputs in place onto the same GPU buffers via session.run(feeds, fetches). Decoding performs zero cache (re-)allocations per step.

Usage

const past_key_values = new StaticCache({ max_cache_len: 4096 });
const output = await generator(messages, { max_new_tokens: 256, past_key_values });
await past_key_values.dispose(); // caller owns the cache

Performance

Decode throughput, q4f16 on WebGPU (TPS), 256 new tokens:

Model	Prefill	DynamicCache	StaticCache	Δ
Phi-4-mini-instruct	2048	10.1	25.4	+151%
Phi-4-mini-instruct	4096	5.6	20.8	+271%
Qwen3.5-4B	2048	17.7	19.2	+8%
Qwen3.5-4B	4096	15.1	18.2	+21%

The gain grows with context length: DynamicCache throughput degrades as buffers get larger while StaticCache stays nearly flat. Qwen3.5 (hybrid attention) benefits less because only its full-attention layers have growing KV entries. In our measurements StaticCache matches a hand-written ORT static-KV baseline within 2%.

Correctness

Token-identical output vs DynamicCache (greedy, 256 new tokens, crossing the 2048/4096 bucket boundaries), on Phi-4-mini-instructand Qwen3.5-4B, Both paths are self-deterministic.

Constraints

Enforced with explicit errors, DynamicCache remains the default everywhere:

• WebGPU only (the problem and the mechanism are WebGPU-specific).
• Decoder-only models, batch size 1.
• The exported graph must support past/present sharing one buffer (past_present_share_buffer semantics — true for GQA-based exports, which current WebGPU decoder exports use). This is why the feature is opt-in.

[nico-martin]

Hi @kylo5aby, thank you so much for looking into this! The performance results are compelling and the overall approach makes sense. Before merging, we need to release per-step tensor wrappers, advance cache length only after successful inference, and make allocation transactional on failure. Could you also add focused lifecycle and WebGPU tests? I'm holding off on merging because the current failure paths can leak resources and corrupt cache state.

[nico-martin]

These are fresh ORT tensor wrappers each run even when they alias the static GPU buffers. Skipping disposal leaves per-step runtime handles to GC. Please explicitly release alias wrappers after each run while preserving the caller-owned buffers, and add a long-generation resource test.

[kylo5aby]

These are fresh ORT tensor wrappers each run even when they alias the static GPU buffers. Skipping disposal leaves per-step runtime handles to GC. Please explicitly release alias wrappers after each run while preserving the caller-owned buffers, and add a long-generation resource test.

Fixed, sessionRun() now reuses the caller's wrappers for pre-allocated outputs, so no per-step alias wrappers exist to leak, and the name-based dispose skip is gone. Covered by the sessionRun identity tests and a 300-step long-generation resource test.

[nico-martin]

This advances sequence length before sessionRun() succeeds. A rejected run leaves the cache logically ahead of its contents and breaks retries. Please commit length only after success or roll back on failure.

[kylo5aby]

This advances sequence length before sessionRun() succeeds. A rejected run leaves the cache logically ahead of its contents and breaks retries. Please commit length only after success or roll back on failure.

Fixed, _reserve() is split into _checkCapacity() (before the run) and _commit() (only after the run succeeded and an in-place-write assertion passed). A failed run leaves the cache length untouched and retryable.

[nico-martin]

If a later allocation throws, earlier buffers/properties remain while allocated is false, so retry allocates over partial state. Please allocate into local structures, commit atomically, and destroy temporary buffers on failure.

[kylo5aby]

If a later allocation throws, earlier buffers/properties remain while allocated is false, so retry allocates over partial state. Please allocate into local structures, commit atomically, and destroy temporary buffers on failure.

Fixed, allocation now goes into local structures and is committed atomically at the end; on failure all created wrappers are disposed and temporary GPU buffers destroyed before rethrowing, so the cache returns to a clean unallocated state and retry works

[kylo5aby]

Hi @kylo5aby, thank you so much for looking into this! The performance results are compelling and the overall approach makes sense. Before merging, we need to release per-step tensor wrappers, advance cache length only after successful inference, and make allocation transactional on failure. Could you also add focused lifecycle and WebGPU tests? I'm holding off on merging because the current failure paths can leak resources and corrupt cache state.

Hi @nico-martin, thanks for the thorough review, all three issues were fixed, plus a few related hardening changes and a dedicated test suite. Added tests/cache_utils.test.js : state-machine lifecycle, transactional allocation with a mocked WebGPU device, generate() integration on a tiny model, sessionRun fetches identity, and the long-generation resource test (300 decode steps: stable wrapper identity, transient outputs disposed, zero GPU buffer churn).

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[nico-martin]

Thanks for the updates! The wrapper cleanup and cache commit ordering look good now, and all 23 focused tests pass locally. There’s just one allocation edge case left: if createBuffer() succeeds but Tensor.fromGpuBuffer() throws, that new buffer never reaches the rollback logic and leaks. Could you destroy it inside createGpuBufferTensor() when wrapping fails and add a test for that path?

[kylo5aby]

Thanks for the updates! The wrapper cleanup and cache commit ordering look good now, and all 23 focused tests pass locally. There’s just one allocation edge case left: if createBuffer() succeeds but Tensor.fromGpuBuffer() throws, that new buffer never reaches the rollback logic and leaks. Could you destroy it inside createGpuBufferTensor() when wrapping fails and add a test for that path?

good catch, that buffer was indeed unreachable by _allocate's rollback since it had not been handed back yet. createGpuBufferTensor() now destroys the buffer and rethrows when Tensor.fromGpuBuffer() fails, and there's a new test covering exactly that path. Also added a "static KV cache" section to the WebGPU guide with usage.

[nico-martin]

Thanks for the updates! LGTM