[justinchuby]

What

The executor's kernel cache is keyed by node and input shapes and never evicted. Every distinct prompt length therefore compiled a fresh kernel for every node in the graph, and each compiled kernel owns device workspaces (attention scratch, dequantization scratch, broadcast metadata). A server whose requests each bring a new prompt length accumulated those workspaces without limit until the device ran out of memory.

This bounds the cache to a small number of shape variants per node, evicting the least recently used beyond it.

Why it looked like something else

The bytes are invisible to the resource governor, which accounts for weights, KV pools and bindings but not for memory a kernel allocates for itself. /v1/resources reported vram.used = weights only while nvidia-smi showed 70+ GiB. That is why the failure first read as "prefill activations scale with prompt length".

What separated the two hypotheses: replaying the same prompt lengths in descending order still grew memory monotonically (43 → 72 GiB). A per-request peak would have shown its maximum on the first, largest request.

Attributing live cuMemAlloc blocks to their allocation sites showed the growth as thousands of retained per-node workspaces (MatMulNBitsKernel::run, GqaWorkspace::reserve, BroadcastMetadataCache::prepare) growing in lockstep at one group per graph node per distinct shape.

Design

• Per node, not global. Eviction stays proportional to what actually varies: a graph keeps its hot prefill and decode variants, and only a node that has genuinely seen many shapes gives one up. A global bound would let one busy node evict every other node's steady-state kernel.
• LRU, not FIFO. The single-token decode shape is the most frequently served and must not be evicted by a run of new prefill shapes.
• Capture safety. The EP's captured device graph is reset before an eviction frees a kernel, because a capture can have baked the pointer of a workspace the eviction is about to release. This mirrors what the device-binding drop path already does.
• Escape hatch. ONNX_RUNTIME_KERNEL_CACHE_VARIANTS_PER_NODE overrides the default of 4.

Measured

30B int4 model, native CUDA backend, 12 requests with prompt lengths from 469 to 5502 tokens:

prompt tokens	before (MiB)	after (MiB)
469	38281	38311
1384	56803	49081
2757	72757	52943
3214	OOM	49829
4587	69915	53585
5045	OOM	50991
5502	69171	47671

Before: three out-of-memory failures. After: every request succeeds, memory is flat, and latency is unchanged (43.7s vs 48.1s on the longest prompt).

Tests

kernel_cache_bounds_the_shape_variants_it_keeps_per_node feeds one node 24 distinct shapes and asserts the cache evicts rather than keeping one entry per shape, while every run still returns the correct result.

cargo test -p onnx-runtime-session (all green), cargo test -p onnx-genai-engine --lib (427 passed), fmt and clippy clean.

Fixes #1362

[github-actions]

⚠️ Benchmark Change Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
⚠️	grammar_masking/llguidance_compute_mask/32	77.02 µs	95.72 µs	+24.3%
⚠️	matmul/medium_generic_f16_threads=8/32x512x512	46.95 µs	57.80 µs	+23.1%
⚠️	qwen3_sampling_processors/top_k_top_p_fast	675.04 µs	824.52 µs	+22.1%
⚠️	logit_processing/seven_processor_chain_per_step	324.59 µs	384.57 µs	+18.5%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.81 ms	6.40 ms	+10.1%
✅	kv_cache/alloc_dealloc_pages	40.10 µs	43.27 µs	+7.9%
✅	qwen3_sampling_processors/top_k_partial_selection	146.11 µs	153.35 µs	+5.0%
✅	gather/small_bf16_threads=1-internal/4096	491.5 ns	515.1 ns	+4.8%
✅	add/large_bf16_threads=1-internal/4194304	1.73 ms	1.80 ms	+3.7%
✅	gather/small_f32_threads=1-internal/4096	659.8 ns	678.2 ns	+2.8%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.07 ms	2.12 ms	+2.6%
✅	matmul/medium_generic_f16_threads=1/32x512x512	38.22 µs	39.02 µs	+2.1%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.19 ms	10.19 ms	+0.0%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	606.75 µs	601.93 µs	-0.8%
✅	reduce_mean/small_f32_threads=1-internal/4096	17.60 µs	17.43 µs	-1.0%
✅	reduce_mean/medium_f32_threads=1-internal/65536	272.08 µs	268.50 µs	-1.3%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	523.86 µs	506.88 µs	-3.2%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.15 ms	2.05 ms	-4.5%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	102.30 µs	95.11 µs	-7.0%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.70 ms	3.43 ms	-7.3%
✅	sampling_latency/top_k_per_token	54.88 µs	49.25 µs	-10.3%
✅	gather/medium_f32_threads=1-internal/32768	4.22 µs	3.76 µs	-10.8%
✅	gather/small_f16_threads=1-internal/4096	545.9 ns	484.0 ns	-11.3%
✅	sampling_latency/top_p_per_token	410.33 µs	362.82 µs	-11.6%
✅	gather/medium_f16_threads=1-internal/32768	2.74 µs	2.42 µs	-11.7%
✅	sampling_latency/min_p_per_token	221.38 µs	193.82 µs	-12.4%
✅	sampling_latency/greedy_per_token	3.43 µs	2.97 µs	-13.4%
✅	add/small_f32_threads=1-internal/1024	234.9 ns	202.9 ns	-13.6%
✅	add/medium_bf16_threads=1-internal/262144	119.44 µs	102.79 µs	-13.9%
🟢	tokenization/encode_tokens_per_second	412.21 µs	349.65 µs	-15.2%
🟢	gather/medium_bf16_threads=1-internal/32768	2.75 µs	2.31 µs	-16.0%
🟢	add/large_f32_threads=1-internal/4194304	802.51 µs	672.18 µs	-16.2%
🟢	matmul/small_generic_f32_threads=1/1x256x256	42.82 µs	35.50 µs	-17.1%
🟢	tokenization/decode_tokens_per_second	7.06 ms	5.78 ms	-18.1%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	6.92 ms	5.57 ms	-19.5%
🟢	reduce_mean/large_f32_threads=1-internal/262144	1.24 ms	973.72 µs	-21.4%
🟢	matmul/small_generic_f16_threads=1/1x256x256	37.07 µs	28.67 µs	-22.7%
🟢	add/medium_f32_threads=1-internal/262144	35.03 µs	26.87 µs	-23.3%
🟢	add/small_f16_threads=1-internal/1024	581.2 ns	444.2 ns	-23.6%
🟢	matmul/small_generic_bf16_threads=1/1x256x256	38.87 µs	28.81 µs	-25.9%
🟢	add/small_bf16_threads=1-internal/1024	636.2 ns	463.1 ns	-27.2%
🟢	add/large_f16_threads=1-internal/4194304	2.31 ms	1.68 ms	-27.5%
🟢	matmul/large_generic_f16_threads=1/32x1024x1024	110.09 µs	79.74 µs	-27.6%
🟢	add/medium_f16_threads=1-internal/262144	144.35 µs	103.30 µs	-28.4%
🟢	matmul/small_generic_f16_threads=8/1x256x256	40.79 µs	29.05 µs	-28.8%
🟢	gather/large_f32_threads=1-internal/131072	43.91 µs	30.53 µs	-30.5%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.87 ms	1.29 ms	-30.8%
🟢	matmul/medium_generic_f32_threads=1/32x512x512	3.16 ms	2.18 ms	-31.1%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	47.15 µs	31.88 µs	-32.4%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.60 ms	1.71 ms	-34.5%
🟢	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	594.41 µs	372.12 µs	-37.4%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	699.37 µs	427.96 µs	-38.8%
🟢	gather/large_bf16_threads=1-internal/131072	18.35 µs	10.63 µs	-42.0%
🟢	matmul/small_generic_f32_threads=8/1x256x256	56.57 µs	32.43 µs	-42.7%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.29 ms	740.52 µs	-42.8%
🟢	gather/large_f16_threads=1-internal/131072	18.19 µs	10.28 µs	-43.5%
🟢	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	96.38 µs	45.98 µs	-52.3%
🟢	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	135.61 µs	48.95 µs	-63.9%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	132.01 µs	45.90 µs	-65.2%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 4.19 3.59 5.21 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)