microsoft / microsoft/onnxruntime

[Web]BucketCacheManager per-bucket limits are far below LLM KV-tensor counts, causing permanent ~3x decode degradation for dynamic-KV workloads

Open
#32,017 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

ep:WebGPU platform:web stale
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 11h
Merged PRs (30d)
184

Description

### Describe the issue

`BUCKET_DEFAULT_LIMIT_TABLE` (`onnxruntime/core/providers/webgpu/buffer_manager.cc`) caps retained buffers per size bucket:

```
{4194304, 20}, {8388608, 10}, {12582912, 10}, {16777216, 10}, {26214400, 15}, ...
```
LLM decode with a growing KV cache (the standard `generate()` pattern: feed `past_key_values`, read `present.*` as `gpu-buffer`) requests **one same-sized buffer per KV output per step** and releases the previous set — 64 buffers/step for a 32-layer model. 64 ≫ cap, so `64 − cap` buffers are truly created (alloc + zero-fill) and destroyed on **every** step.

Result: a permanent per-token latency staircase, steps exactly on bucket boundaries. Phi-4-mini q4f16 (64 KV tensors, KV tensor = `2048·S` bytes), onnxruntime-web 1.29, Intel iGPU:

| S | bucket | cap | miss | ms/token |
| --- | --- | --- | --- | --- |
| < 2048 | 4 MiB | 20 | 44 | ~46 |
| 2048–4096 | 8 MiB | 10 | 54 | ~138 |
| 4096–6144 | 12 MiB | 10 | 54 | ~233 |
| > 6144 | 16 MiB | 10 | 54 | ~298 |

Decode drops ~22 → ~4 TPS over 8k context.

Evidence pinning this to the bucket table:

- Step at S=6144 = the 12→16 MiB boundary; 12 MiB is not a power of two, so generic size-class effects predict nothing there.
- Qwen3.5-4B (16 growing KV tensors, fits under caps): no comparable staircase, and the predicted *negative* step at 16→25 MiB (cap 10→15, miss 6→1) measured as −6.0 ms.
- - Zero-churn control (pre-allocated KV bound in-place via fetches, same graph): no staircase, ~3x faster at long context.

### Why existing options don't cover this

- `storageBufferCacheMode: 'disabled'`: destroys released buffers immediately → use-after-destroy validation error with GPU-resident KV I/O binding.
- `'simple'`: caches by exact size, never evicts → dynamic shapes miss every step *and* the pool grows unboundedly → device lost. (Fine for the static-shape case in #29016 — different workload.)
- `'bucket'` (default): works, but fixed caps produce the above.
- #29017 exposes mode selection to JS; no mode handles "N same-sized buffers per step, N > cap, size drifting across buckets".

### Proposed direction

Simply raising the caps doesn't work: the needed cap is model-dependent (40 layers → 80 tensors), and decode walks through every bucket over time, so cap=64 puts worst-case retained memory near ~4 GiB (vs ~815 MiB today). Alternatives:

1. Adapt per-bucket caps to the observed acquire rate (with decay), or
2. Cross-bucket eviction under a total-bytes budget — the bucket a growing KV has moved out of is pure dead weight today, or
3. Expose the limit table via EP options (following #29017), letting callers size it to their KV tensor count.

### To reproduce

In Transformers.js/ort-web:

1. Load Phi-4-mini-instruct ONNX q4f16, device: "webgpu".
2. Prompt ~1900 tokens, generate 256 greedily, timestamp each TextStreamer.put.
3. Per-token latency jumps ~46 → ~138 ms permanently at position ~2050. Prefill ~3900 / ~6000 shows the 4096 / 6144 steps.

### Urgency

_No response_

### ONNX Runtime Installation

Built from Source

### ONNX Runtime Version or Commit ID

1.29.0(commit id: 7ad118c)

### Execution Provider

'webgpu' (WebGPU)

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in onnxruntime/core/providers/webgpu/buffer_manager.cc by reading BUCKET_DEFAULT_LIMIT_TABLE and the bucket cache acquire/release behavior. Reproduce the staircase with Phi-4-mini q4f16 in Transformers.js/ort-web using the WebGPU provider, then evaluate a bounded approach against the reported dynamic-KV workload. Done means the per-token staircase is reduced without unbounded pool growth, excessive retained memory, or device loss.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.