Skip to content

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

Open
kylo5aby wants to merge 2 commits into
huggingface:mainfrom
kylo5aby:feat/static-cache
Open

Add StaticCache: fixed-capacity in-place KV cache for WebGPU#1745
kylo5aby wants to merge 2 commits into
huggingface:mainfrom
kylo5aby:feat/static-cache

Conversation

@kylo5aby

Copy link
Copy Markdown

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.

Signed-off-by: Zhenwei Jin <zhenwei.jin@intel.com>
Signed-off-by: Zhenwei Jin <zhenwei.jin@intel.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[WebGPU] WebGPU decode permanently degrades ~3x at S=2048/4096/6144

1 participant