[Feature]: Layer-wise KV Cache Async Transfer to Reduce TTFT in Disaggregated Serving

🚀 The feature, motivation and pitch

Feature
Implement a layer-wise KV cache async transfer mechanism for TensorRT-LLM in disaggregated serving scenarios.

Motivation

The current disaggregated serving implementation requires a full KV cache transfer after prefill computation, leading to TTFT = prefill computation time + full KV transfer time. This makes TTFT for long contexts significantly slower than in frameworks like vLLm, which supports layer-wise KV async transfer.

The key advantage of layer-wise KV async transfer is its ability to overlap communication with computation. By transferring each layer's KV cache as soon as it's computed, the transmission of early layers is hidden behind the computation of subsequent ones. This makes the transfer process largely asynchronous. As a result, the effective TTFT is drastically reduced to just the prefill computation time plus the transfer time of the final layer's KV cache:

TTFT ≈ prefill computation time + last layer KV transfer time

This change is critical for achieving low-latency first token generation, especially with long context windows.

Alternatives

vLLM with LMCache.
https://docs.lmcache.ai/kv_cache_optimizations/layerwise.html

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.