混合#
CacheBlend 通过在非前缀位置重计算一部分令牌来实现 KV Cache 的重用。例如,当相应的文本在 LLM 输入中连接时,CacheBlend 可以组合多个(预)计算的 KV Cache。
在 RAG 场景中配置 CacheBlend#
在这里,我们将解释我们端到端的 example 中的代码。
以下是一些与混合相关的配置(及其解释):
# Enable blending in LMCache
os.environ["LMCACHE_ENABLE_BLENDING"] = "True"
# Separator string between different chunks
os.environ["LMCACHE_BLEND_SPECIAL_STR"] = " # # "
# Layerwise must be turned on when blending is enabled
os.environ["LMCACHE_USE_LAYERWISE"] = "True"
# Determining which tokens to recompute at layer 1
os.environ["LMCACHE_BLEND_CHECK_LAYERS"] = "1"
# Ratio of tokens to recompute
os.environ["LMCACHE_BLEND_RECOMPUTE_RATIOS"] = "0.15"
# Optionally, we can use sparse attention to improve generation quality
# by using more accurate attention mask
if enable_sparse:
os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER"
os.environ["LMCACHE_EXTRA_CONFIG"] = '{"enable_sparse": true}'
首先,我们将文本预处理为标记,因为将连接的字符串进行标记化可能会产生与分别对每个字符串进行标记化后再连接的结果不同的标记。例如,假设我们有一个系统提示和三个文本块。在将它们发送到 LLM 之前,我们需要将它们预处理为标记:
sys_prompt = tokenizer.encode("You are a very helpful assistant.")
chunk1_prompt = tokenizer.encode("Hello, how are you?" * 500)[1:]
chunk2_prompt = tokenizer.encode("Hello, what's up?" * 500)[1:]
chunk3_prompt = tokenizer.encode("Hi, what are you up to?" * 500)[1:]
blend_special_str = tokenizer.encode(os.getenv("LMCACHE_BLEND_SPECIAL_STR"))[1:]
first_prompt = (
sys_prompt
+ blend_special_str
+ chunk1_prompt
+ blend_special_str
+ chunk2_prompt
+ blend_special_str
+ chunk3_prompt
+ blend_special_str
+ tokenizer.encode("Hello, my name is")[1:]
)
然后,我们可以将标记化的提示发送到 vLLM。同时,LMCache 将根据 BLEND_SPECIAL_STR 存储不同块的 KV 缓存。
llm.generate(prompts={"prompt_token_ids": first_prompt})
同样,我们使用相同的块但以不同的顺序构建另一个提示。
second_prompt = (
sys_prompt
+ blend_special_str
+ chunk2_prompt
+ blend_special_str
+ chunk1_prompt
+ blend_special_str
+ chunk3_prompt
+ blend_special_str
+ tokenizer.encode("Hello, how are you?")[1:]
)
llm.generate(prompts={"prompt_token_ids": second_prompt})
尽管第二个提示的块顺序不同,LMCache 仍然可以重用 chunk1、chunk2 和 chunk3 的 KV 缓存。