[bvolpato]

Description

This change fixes a GPU memory leak in the optimized KDA prefill op (tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py). Without the change, a server whose prefill length changes on almost every call grew from 142.8 GiB to 178.2 GiB of GPU memory in use per GPU in 20 minutes. Then it stopped with CUDA out of memory. With the change, a 44-minute run finished 897 calls with 0 failed calls, and the memory in use stayed at 144 GiB.

Pull request 19724 of this repository is older, is open, and fixes the same leak in another way (it keeps the wrappers inside the _buf_cache entry), so this draft may be closed in its favor.

Cause. _get_buffers keeps the scratch tensors of the last 8 prefill shapes in _buf_cache (LRU). _launch_fused_k123_inv wraps seven of those tensors with _ct_cached. The wrappers live in _input_wrap_cache and pin the tensor storage. _prune_on_gc removes a wrapper only when its tensor is garbage-collected, and that cannot happen while the wrapper pins the tensor. So the seven tensors of every evicted entry stay allocated. A benchmark with a fixed input length never evicts an entry. A serving process with a new prefill length on almost every call evicts entries all the time.

Change. _get_buffers now drops the _ct_cached wrappers of an entry when _buf_cache evicts it (new helper _release_scratch_wrappers).

Measurement. One server on 4x B200 (TP4, EP4, MTP) with an FP8 GLM-5.3-Flash checkpoint and an agent workload: multi-turn sessions with prompts near 100K tokens, where almost every call has a new prefill length. The memory values are GPU memory in use per GPU.

Run	Server settings	Client load	GPU memory in use	Result
Without the change	free_gpu_memory_fraction: 0.6, --max_num_tokens 8192, MTP draft length 3	16 requests in flight	142.8 GiB to 178.2 GiB in 20 minutes (about 1.75 GiB per minute)	CUDA out of memory in the sparse prefill attention after 485 calls
Without the change	free_gpu_memory_fraction: 0.5, --max_num_tokens 16384, MTP draft length 4 with rejection sampling	4 requests in flight	136.9 GiB to 171.6 GiB in 21 minutes	293 calls, 0 failed. The same server later ran out of memory at 16 requests in flight
With the change	Same as the row above	16 requests in flight	144 GiB	44 minutes, 897 calls, 0 failed

No run without the change started with both the server settings and the client load of the run with the change. The second row has the same settings at a lower load. That server got 16 requests in flight only after the growth of the second row, and I have no memory values for that part.

In the first run, the KV cache blocks in use stayed between 15K and 34K of 81,688 during the growth, so the KV cache did not fill up. A tensor census in a debug run showed BF16 scratch tensors of shape [1, T, 16, 128] left behind for each distinct prefill length T. An earlier run with free_gpu_memory_fraction: 0.75 and --max_num_tokens 16384 stopped with the same error after 6 minutes and 204 calls.

Expected improvement and risk

Change	Expected improvement	Risk	What limits or turns it off
_get_buffers releases the _ct_cached wrappers of each entry that _buf_cache evicts	GPU memory only. Condition: 16 requests in flight and a new prefill length on almost every call. With the change, a 44-minute run of 897 calls had 0 failed calls and 144 GiB in use per GPU. Without it, the memory in use grew by about 1.75 GiB per minute and the server stopped with CUDA out of memory after 20 minutes (other server settings, see the table above). I expect no change in output and no measurable change in speed. I measured neither.	Low. The new code runs only when _buf_cache evicts an entry. The comment on _buf_cache says that the LRU bound exists so that scratch is not pinned forever. The wrapper dict inside an entry already leaves with the entry. This change does the same for the wrappers that _ct_cached keeps outside the entry.	No flag turns it off. It runs only at an eviction, so a process with 8 or fewer distinct prefill shapes (_BUF_CACHE_MAX_ENTRIES) never runs it. The cost is one scan of _input_wrap_cache for each eviction.

Test Coverage

• New tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py::test_evicted_prefill_scratch_is_released. The test runs the indexed prefill with more distinct lengths than the buffer cache holds. Then it checks that no cached wrapper pins an evicted scratch tensor. On B200 it fails without the change ("28 evicted scratch tensors are still pinned by cached wrappers", which is 4 evictions of 7 tensors each) and passes with the change.
• The test file is in the B200 pre-merge list tests/integration/test_lists/test-db/l0_b200.yml.
• The serving runs of the table used the 1.3.0rc29 wheel with the Python files of [None][feat] Add GLM-5.3-Flash (glm5_next) support #19136 and [None][fix] wait for K1 readers before reusing KDA prefill stages #19581 and the code change of [None][fix] Use draft slot tables for GLM-5.3-Flash MTP #19803 on top. The run with the change also had the code change of this pull request as a local patch (the same statements, other docstring text). The op file at the release tag of that wheel (v1.3.0rc29) is the same as the file at the base of this pull request. I checked the output of the serving run only lightly: answers are sensible and tool calls parse.

To verify, run the new test on a Blackwell GPU (SM100 or SM103) with the fla package installed:

pytest tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py::test_evicted_prefill_scratch_is_released

To see the leak in serving, send prompts whose prefill length changes on every call to a model that uses the optimized KDA prefill, and watch the GPU memory in use.

Not verified

• Speed. I have no latency or throughput comparison with and without the change.
• Memory over time with the change. My notes have 144 GiB for the 44-minute run and 140.9 GiB for an earlier check of the same setup. I did not verify that the two values come from one server process, so this text gives no growth rate.
• A like-for-like run without the change: a new server with the settings and the client load of the run with the change.
• Freed GPU memory in the test. The new test checks the wrapper cache. It does not measure the memory that an eviction frees.
• Kernel output after an evict and refill cycle, and an eviction while a kernel on another CUDA stream can still read the scratch.
• The other tests of tests/unittest/_torch/modules/kimi_kda/ and the other GPU and integration tests of the repository. No upstream CI has run on this pull request.
• A wheel built from this branch.
• Hardware other than B200.

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.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

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

• The reviewers assigned automatically/manually are appropriate for the PR.

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

🤖 Generated with Claude Code

[longlee0622]

Superseded by #19854