From 7f797780685c6d8997748b7e9781e17e8011c3ac Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 04:04:33 +0000 Subject: [PATCH] Bound the shape-keyed kernel cache per node (#1362) The executor's kernel cache is keyed by node *and* input shapes and never evicted, so every distinct prompt length compiled a fresh kernel for every node in the graph. A compiled kernel owns device workspaces (attention scratch, dequantization scratch, broadcast metadata), so a server whose requests each bring a new prompt length accumulated those workspaces without limit until the device ran out of memory. The bytes were invisible to the resource governor, which accounts for weights, KV pools, and bindings but not for memory a kernel allocates for itself. That is why the failure read as "activations scale with prompt length": VRAM grew with the number of *distinct* shapes seen, not with the size of any one request. Replaying the same prompt lengths in descending order still grew memory monotonically, which is what separated an accumulating cache from a per-request peak. Keep at most a bounded number of shape variants per node, evicting the least recently used beyond it. The bound is per node rather than global so 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. 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. Measured on a 30B int4 model on the native CUDA backend, 12 requests with prompt lengths from 469 to 5502 tokens: device memory climbed 38 GiB to 72.7 GiB with three out-of-memory failures, and now stays flat at 45-53 GiB with every request succeeding and no change in latency. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 0190e2eb-abe4-451f-b36d-44a035a99b7e --- .../src/executor/kernel_cache.rs | 93 +++++++++++++++++++ crates/onnx-runtime-session/tests/executor.rs | 47 ++++++++++ 2 files changed, 140 insertions(+) diff --git a/crates/onnx-runtime-session/src/executor/kernel_cache.rs b/crates/onnx-runtime-session/src/executor/kernel_cache.rs index b12357ac13..c928b69cd2 100644 --- a/crates/onnx-runtime-session/src/executor/kernel_cache.rs +++ b/crates/onnx-runtime-session/src/executor/kernel_cache.rs @@ -626,24 +626,112 @@ pub struct CacheStats { pub misses: u64, /// Lookups served via the pre-bound fast path (zero-alloc). pub prebind_hits: u64, + /// Entries dropped by the per-node variant bound (issue #1362). + pub evictions: u64, } /// Shape-keyed kernel cache (§11.1). Owns the compiled kernels for the session. #[derive(Default)] pub(crate) struct KernelCache { pub(super) entries: HashMap>, + /// Logical time each entry was last served, for the per-node bound below. + /// Kept beside `entries` rather than inside them so a shared-reference hit + /// can record its use without the map itself becoming interior-mutable. + pub(super) last_used: HashMap, + /// Monotonic tick handed out to `last_used`. + pub(super) clock: AtomicU64, pub(super) hits: u64, pub(super) misses: u64, + /// Entries dropped by the per-node bound. + pub(super) evictions: u64, pub(super) prebind_hits: AtomicU64, } +/// How many shape variants of one node the cache keeps (issue #1362). +/// +/// The cache is keyed by node **and input shapes**, so every distinct prompt +/// length compiles a fresh kernel for every node in the graph. A compiled kernel +/// owns device workspaces (attention scratch, dequantization scratch, broadcast +/// metadata), so an unbounded cache turns "each request has a new prompt length" +/// into VRAM that grows without limit until the device is exhausted — memory the +/// resource governor never sees, because it belongs to the kernels rather than +/// to any binding or KV pool. +/// +/// A per-node bound rather than a global one keeps the eviction proportional to +/// what actually varies: a graph keeps its hot decode and prefill variants, and +/// only a node that has genuinely seen many shapes gives one up. The default is +/// deliberately generous enough to hold the shapes one request cycles through +/// (a chunked prefill shape, its remainder, and the single-token decode shape) +/// so steady state never recompiles. +const DEFAULT_VARIANTS_PER_NODE: usize = 4; + +fn variants_per_node() -> usize { + static RESOLVED: std::sync::OnceLock = std::sync::OnceLock::new(); + *RESOLVED.get_or_init(|| { + std::env::var("ONNX_RUNTIME_KERNEL_CACHE_VARIANTS_PER_NODE") + .ok() + .and_then(|value| value.trim().parse::().ok()) + .filter(|variants| *variants > 0) + .unwrap_or(DEFAULT_VARIANTS_PER_NODE) + }) +} + impl KernelCache { + /// Next logical tick. + fn tick(&self) -> u64 { + self.clock.fetch_add(1, Ordering::Relaxed) + } + + /// Record that `key` was just served, so eviction can tell the hot variants + /// of a node from the ones a single past prompt length left behind. + fn touch(&self, key: &KernelKey) { + if let Some(slot) = self.last_used.get(key) { + slot.store(self.tick(), Ordering::Relaxed); + } + } + + /// Drop the least recently used variants of `node` beyond the per-node bound. + /// + /// A kernel's device workspaces are freed by its `Drop`, so the EP's captured + /// device graph is reset first: a capture can have baked the pointer of a + /// workspace this eviction is about to free, and replaying it afterwards + /// would read freed device memory. Resetting is what the device-binding drop + /// path already does for the same reason. + fn evict_surplus_variants(&mut self, node: u32, ep: &dyn ExecutionProvider) { + let bound = variants_per_node(); + let mut variants = self + .entries + .keys() + .filter(|key| key.node == node) + .map(|key| { + let used = self + .last_used + .get(key) + .map(|slot| slot.load(Ordering::Relaxed)) + .unwrap_or(0); + (used, key.clone()) + }) + .collect::>(); + if variants.len() <= bound { + return; + } + variants.sort_by_key(|(used, _)| *used); + let surplus = variants.len() - bound; + let _ = ep.reset_device_graph(); + for (_, key) in variants.into_iter().take(surplus) { + self.entries.remove(&key); + self.last_used.remove(&key); + self.evictions += 1; + } + } + pub(super) fn stats(&self) -> CacheStats { CacheStats { entries: self.entries.len(), hits: self.hits, misses: self.misses, prebind_hits: self.prebind_hits.load(Ordering::Relaxed), + evictions: self.evictions, } } @@ -665,6 +753,7 @@ impl KernelCache { return None; } let kernel = self.entries.get(binding)?.as_ref(); + self.touch(binding); self.prebind_hits.fetch_add(1, Ordering::Relaxed); #[cfg(test)] PREBIND_FAST_PATH_TEST_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); @@ -739,8 +828,12 @@ impl KernelCache { kernel.set_constant_inputs(constant_inputs); kernel.set_capture_seq_independent(capture_seq_independent); self.entries.insert(key.clone(), kernel); + self.last_used + .insert(key.clone(), AtomicU64::new(self.tick())); self.misses += 1; + self.evict_surplus_variants(key.node, ep); } + self.touch(&key); #[cfg(test)] PREBIND_FALLBACK_TEST_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); let kernel_ref = self.entries.get(&key).expect("just inserted").as_ref(); diff --git a/crates/onnx-runtime-session/tests/executor.rs b/crates/onnx-runtime-session/tests/executor.rs index 13c3ac3e0c..b83d74a1dd 100644 --- a/crates/onnx-runtime-session/tests/executor.rs +++ b/crates/onnx-runtime-session/tests/executor.rs @@ -1337,6 +1337,53 @@ fn symbolic_batch_matmul_chain_runs_for_multiple_shapes() { assert_eq!(s2b.hits, 1, "the fused node served from the batch=2 cache"); } +/// A shape-keyed kernel cache that never evicts turns "every request has a new +/// prompt length" into unbounded device memory, because each compiled kernel +/// owns its own workspaces (issue #1362). Feeding one node many distinct shapes +/// must therefore settle at the per-node bound instead of growing with the +/// number of distinct shapes seen. +#[test] +fn kernel_cache_bounds_the_shape_variants_it_keeps_per_node() { + let w_data = [ + 1.0f32, 0.0, 0.0, 0.0, // + 0.0, 1.0, 0.0, 0.0, // + 0.0, 0.0, 1.0, 0.0, // + 0.0, 0.0, 0.0, 1.0, + ]; + + let mut g = Graph::new(); + let batch = g.intern_symbol("batch"); + let sym_row = || vec![Dim::Symbolic(batch), Dim::Static(4)]; + let x = input_shaped(&mut g, "X", DataType::Float32, sym_row()); + let w = f32_init(&mut g, "W", &[4, 4], &w_data); + let m = op_shaped(&mut g, "MatMul", &[x, w], DataType::Float32, sym_row(), &[]); + g.add_output(m); + + let mut session = InferenceSession::from_graph(g).expect("build symbolic session"); + + // Far more distinct shapes than any per-node bound would keep. + let shapes = 24; + for rows in 1..=shapes { + let data: Vec = (0..rows * 4).map(|i| i as f32).collect(); + let x_tensor = Tensor::from_f32(&[rows, 4], &data).unwrap(); + let out = session.run(&[("X", &x_tensor)]).expect("run"); + assert_eq!(out[0].shape, vec![rows, 4]); + // Identity weights: the output must still be the input, so eviction is + // shown to cost recompilation and nothing else. + assert_close(&out[0].to_vec_f32(), &data); + } + + let stats = session.cache_stats(); + assert!( + stats.evictions > 0, + "the bound must have evicted surplus variants, saw {stats:?}" + ); + assert!( + stats.entries < shapes, + "the cache must not keep one entry per distinct shape, saw {stats:?}" + ); +} + /// Two inputs share a symbol (`batch`); supplying them with *conflicting* /// concrete sizes is a resolution error, not a silently-wrong run. #[test]