[Feature]: cache transfer without extra GPU memory allocation for long ISL requests

🚀 The feature, motivation and pitch

I tried to deploy LLMs for agentic usage with PD disaggregation on UCX backend. I set max_tokens_in_buffer of transciever config to 10240 first. But in the real scenario, the max input length could easily exceed 10240, and I observed that huge GPU memory is allocated accumulatively until OOM is raised (The relative codes I think are in BaseTransBufferManager::getOrAllocateBuffers). However, if I tried to increase max_tokens_in_buffer of transciever config as recommended in docs, memory always run out during server launch because the pre-allocated memory is too large, and memory fraction should be reduced to an extremely low value.

I also tried to WAR with layer-wise or zero copy methods supplied.

First, when I set TRTLLM_DISAGG_LAYERWISE to 1, Error 'Window Size 0 not found' was raised in CacheFormatter::unformat, and after I fixed the bug, illegal memory access error was raised in concatKVCacheDispatch kernel.

Then I tried the zero copy method with TRTLLM_TRY_ZCOPY_FOR_KVCACHE_TRANSFER set as 1, but the program hung during cache transfer. I found the bug in the conditions whether zero copy kv cache transfer is allowed to be executed in CacheFormatter::unformat and CacheFormatter::format functions, which should be the same but actually not. After I fixed it, zero copy could work. But considering its strict restrictions on parallelism configs, the way of zero copy does not help for real scenarios.

In my opinion, it is crucial to solve the problem of exaggerated GPU memory allocation for cache transfer if we want to deploy the LLMs for real scenarios. I'm looking forward to an improved cache transfer method for long ISL (such as 128k) for flexible parallelism configs. If this feature has already been supported, please tell me how to enable it. Thank you!

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.