[vsabavat]

Description

The Kimi K3 KDA prefill op (tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py) caches its intermediate
buffers per token-count shape in _buf_cache. That is an LRU of at most _BUF_CACHE_MAX_ENTRIES (8) entries; a new
shape evicts the oldest entry. In serving, most prefill batches bring a new token count, so entries are evicted all
the time.

The K123 and akk_inv launches wrapped seven of those buffers (k_scaled, kg, q_scaled, beta_activated,
gk_last_exp, A_qk, A_kk) with _ct_cached.

• _ct_cached kept each CuTe wrapper in a module-level dict keyed by id(tensor), and relied on a
  weakref.finalize on the tensor to drop the wrapper.
• But a wrapper made by from_dlpack holds a DLPack capsule, which owns a reference to the tensor. The tensor is never
  collected, so the finalizer never runs.
• So every evicted entry's buffers stay allocated for the life of the process. This was introduced with the KDA prefill
  scratch cache.

Impact: GPU memory grows without bound whenever prefill shapes vary, by about 1 KiB per token per KDA head for every
evicted shape.

• For Kimi K3 at TP16 (6 KDA heads per rank), with GSM8K-length prefill batches, that is about 9 MiB per evicted shape
  per rank.
• Serving GSM8K on one LLM instance grew allocated memory by about 0.5 GiB per 200 requests on every rank.
• gc.collect() reclaims none of it.

The fix:

• _scratch_ct caches each buffer's wrapper in the cute_wrappers dict of that buffer's own _buf_cache entry, so
  evicting the entry releases the buffers and their wrappers together.
• A cache hit still reuses the wrappers, with no per-call from_dlpack.
• _ct_cached and its dict have no other users and are removed.

This PR and #19835 (a new prefill output tensor per call) change different lines of the module and apply in either
order.

Test Coverage

tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py (in l0_b200.yml):

• test_evicted_prefill_scratch_is_released (one GPU): runs indexed varlen prefill in cyclic passes over 12 token
  counts. The cache holds 8, so every call misses and evicts.
  • After every pass but the first, it checks that the cache holds at most 8 entries and that
    torch.cuda.memory_allocated() is within 16 MiB of its value after the first pass.
  • It also checks that the buffers of an entry cached after the first pass are freed once the next pass evicts it
    (their weak references are dead).
  • Fails before the fix: about 1 GiB more allocated after one pass.

Results

One GB200 GPU, the whole test file, with the module as on main at cf5dc9f and with this PR:

• Before: the new test fails ("pass 0: 1050.8 MiB more allocated after the evictions"); 3 passed, 2 skipped.
• With this PR: 4 passed, 2 skipped.
• The 2 skips are test_recycled_cu_seqlens_id_matches_fla ("could not obtain a recycled cu_seqlens tensor ID"),
  the same before and after.

End to end, Kimi K3 at TP16 (16 GB200 GPUs) serving GSM8K-200 five times on one LLM instance:

• With this fix, torch.cuda.memory_allocated() on every rank stays within 0.3 GiB of its round-1 value over the five
  rounds, with no trend (the minimum over ranks is 112,562-112,569 MiB).
• Without it, the same run grew by about 0.5 GiB per round on every rank.
• Accuracy is the same with and without the fix: 96.0 / 98.0 / 97.0 / 96.5 / 97.0 over the five rounds.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title. (No API change.)

• Any new dependencies have been scanned for license and vulnerabilities (No new dependencies.)

• CODEOWNERS updated if ownership changes (No ownership change.)

• Documentation updated as needed (No user-facing change.)

• Update tava architecture diagram if there is a significant design change in PR. (No design change.)

• The reviewers assigned automatically/manually are appropriate for the PR. (To check once the PR is open.)

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[vsabavat]

Superseded by the consolidated draft PR: #19988. Source commits and follow-ups are retained; prior reviews remain historical evidence, not approval of the replacement heads.