Bound the shape-keyed kernel cache per node (#1362)

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<KernelKey, Box<dyn onnx_runtime_ep_api::Kernel>>,
+    /// 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<KernelKey, AtomicU64>,
+    /// 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<usize> = 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::<usize>().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::<Vec<_>>();
+        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();

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<f32> = (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]