Can we cache some pre-known prompt word sets K and V for use in the Context phase?

[yifeihappy]

In my application scenario, I hope to know whether the product is related to our search terms; the collection of products is known in advance. How can it be optimized? My idea is to cache the K and V related to the title of the product first, and then concatenate them with the search terms in the Context phase and use them directly. Based on TensorRT-llm, can this goal be achieved?

[yifeihappy]

I very much hope that tensor-llm can open source the implementation of batch manager, because there are many functional implementation details for reference.

[schetlur-nv]

@yifeihappy we have an experimental feature to reuse KV cache blocks across requests. You can enable it with https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/batch_manager/kvCacheConfig.h#L40
With this feature, if the first N tokens of two requests are the same, the KV cache values for those tokens will get computed only once and automatically reused.
There are some limitations like the reuse happens at the KV cache block granularity, and so you need sufficient shared tokens to fill at least one KV cache block to get this optimization. Today that value is 128 tokens.
We are in the process of making this more general (feature on by default and <128 tokens per KV cache block) soon.

[yifeihappy]

@yifeihappy we have an experimental feature to reuse KV cache blocks across requests. You can enable it with https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/batch_manager/kvCacheConfig.h#L40 With this feature, if the first N tokens of two requests are the same, the KV cache values for those tokens will get computed only once and automatically reused. There are some limitations like the reuse happens at the KV cache block granularity, and so you need sufficient shared tokens to fill at least one KV cache block to get this optimization. Today that value is 128 tokens. We are in the process of making this more general (feature on by default and <128 tokens per KV cache block) soon.

"enableBlockReuse" seems to be used in GptManager. Because this part of the code is closed source and I can't see the usage details inside, I have a few questions about enableBlockReuse:

1. Do two requests that can share the first N tokens need to be in the same batch?
2. Is it possible to pre-calculate the K and V of some popular prompts and cache them when the service starts?
3. Because the video memory size is limited, is there a cache management strategy, such as LRU, to update and eliminate the K and V of pre-calculated prompts?

[schetlur-nv]

1. No. They can arrive at different points of time.
2. You can do this by sending requests with popular system prompts in a warmup phase for the server. LRU eviction policy will naturally make the system prompt persist in GPU memory for longer as long as requests are using it. If system prompt stops being used, it will eventually be evicted.
3. Yes, we use an LRU strategy. We slightly modify this to preferentially free the leaf tokens (i.e. last tokens) of different requests that all have the same idle time first. You can think of it as a "breadth-first" LRU

[nv-guomingz]

hi do u still have further issue or question now? If not, we'll close it soon.