[None][fix] Release the cute wrappers of evicted KDA prefill scratch

tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py

@@ -187,16 +187,29 @@ def _prune_on_gc(cache, key, *keyobjs):
         weakref.finalize(o, cache.pop, key, None)
 
 
+def _release_scratch_wrappers(entry):
+    """Drop the cached cute wrappers of scratch tensors that leave an LRU cache.
+
+    A wrapper from ``_ct_cached`` pins its tensor's storage, so the finalizer
+    registered by ``_prune_on_gc`` never runs for it. Evicted scratch must
+    release its wrappers here, or every evicted shape stays allocated forever.
+    """
+    ids = {id(t) for t in entry if isinstance(t, torch.Tensor)}
+    for key in [k for k in _input_wrap_cache if k[0] in ids]:
+        _input_wrap_cache.pop(key, None)
+
+
 def _ct_cached(t, etype):
     """`_ct(t, etype)` with id(t)-based cache. Returns the same cute wrapper
     for repeated calls with the same tensor object, avoiding per-call
     `from_dlpack` overhead (~5-10us each).
 
-    ONLY use for tensors with process-long lifetime (module params, the
-    module-level scratch from ``_get_buffers``): the cached wrapper pins the
-    tensor's storage, so the weakref pruning never fires for the keyed
+    ONLY use for tensors with process-long lifetime (module params) and for
+    the module-level scratch from ``_get_buffers``: the cached wrapper pins
+    the tensor's storage, so the weakref pruning never fires for the keyed
     object and a per-call activation would be pinned forever (~100MB/call
     leak in the executor runtime). Per-call tensors must use plain ``_ct``.
+    ``_get_buffers`` releases the wrappers of the scratch it evicts.
     """
     key = (id(t), etype)
     w = _input_wrap_cache.get(key)
@@ -471,7 +484,7 @@ def _wrap(t, etype):
         )
 
         while len(_buf_cache) >= _BUF_CACHE_MAX_ENTRIES:
-            _buf_cache.pop(next(iter(_buf_cache)))
+            _release_scratch_wrappers(_buf_cache.pop(next(iter(_buf_cache))))
         _buf_cache[key] = (
             k_scaled,
             kg,

tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py

@@ -246,6 +246,43 @@ def test_recycled_cu_seqlens_id_matches_fla(
     assert module._varlen_pure_cache.get(id(cu_seqlens)) is False
 
 
+@torch.no_grad()
+def test_evicted_prefill_scratch_is_released(
+    dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],
+    gate_params: tuple[torch.Tensor, torch.Tensor],
+) -> None:
+    """Scratch that leaves the LRU buffer cache keeps no pinned cute wrapper.
+
+    A serving process sees a new prefill length on almost every call. The
+    buffer cache holds the scratch of the last few lengths only, so the
+    scratch of every older length must be freed.
+    """
+    optimized, _ = dispatch_pair
+    module = _op_module()
+
+    def scratch_ids() -> set[int]:
+        return {
+            id(tensor)
+            for entry in module._buf_cache.values()
+            for tensor in entry
+            if isinstance(tensor, torch.Tensor)
+        }
+
+    seen: set[int] = set()
+    for step in range(module._BUF_CACHE_MAX_ENTRIES + 4):
+        total_tokens = 256 + 64 * step
+        inputs = _make_inputs(total_tokens, seed=100 + step)
+        run_indexed_prefill(optimized, gate_params, inputs, _make_cu_seqlens([total_tokens]))
+        torch.cuda.synchronize()
+        seen |= scratch_ids()
+
+    assert len(module._buf_cache) <= module._BUF_CACHE_MAX_ENTRIES
+    evicted = seen - scratch_ids()
+    assert evicted, "the run must evict at least one buffer-cache entry"
+    pinned = {key[0] for key in module._input_wrap_cache} & evicted
+    assert not pinned, f"{len(pinned)} evicted scratch tensors are still pinned by cached wrappers"
+
+
 @torch.no_grad()
 def test_repeated_single_sequence_metadata_matches_fla(
     dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch],