chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,88 @@
|
||||
Blending
|
||||
================
|
||||
|
||||
.. warning::
|
||||
|
||||
This page documents the behavior of LMCache's in-process mode (deprecated). Please consider using :doc:`LMCache MP mode </mp/index>` for better feature support and performance.
|
||||
|
||||
|
||||
CacheBlend enables KV cache reuse for non-prefix positions by recomputing a subset of tokens at non-prefix positions.
|
||||
For example, CacheBlend can combine multiple (pre-)computed KV caches, when their corresponding texts are concatenated in the LLM input
|
||||
|
||||
Configuring CacheBlend in RAG scenarios
|
||||
-------------------------------------------------
|
||||
|
||||
Here, we will explain the code in our end-to-end `example <https://github.com/LMCache/LMCache/tree/dev/examples/blend_kv_v1/blend.py>`_>.
|
||||
|
||||
Below are some blending-related configurations (and explanations):
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# 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}'
|
||||
|
||||
Firstly, we preprocess texts into tokens, as tokenizing a concatenated string may produce different tokens than concatenating the results of tokenizing each string individually.
|
||||
For example, assume we have a system prompt and three text chunks. We need to preprocess them into tokens before sending to the LLM:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
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:]
|
||||
)
|
||||
|
||||
Then, we can send the tokenized prompt to vLLM. Meanwhile, LMCache will store the KV caches of different chunks according to the ``BLEND_SPECIAL_STR``.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
llm.generate(prompts={"prompt_token_ids": first_prompt})
|
||||
|
||||
Similarly, we build another prompt using the same chunks but with different orders.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
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})
|
||||
|
||||
Even though the second prompt has a different order of chunks, LMCache can still reuse the KV caches of chunk1, chunk2, and chunk3.
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
CacheBlend
|
||||
==========
|
||||
|
||||
CacheBlend lets LMCache reuse the KV cache of **any** repeated text chunk --
|
||||
not only a shared prefix -- by selectively recomputing a small fraction of
|
||||
tokens at chunk boundaries. This cuts time-to-first-token for RAG and
|
||||
multi-document workloads where the reusable context is not a clean prefix.
|
||||
|
||||
Enabling CacheBlend (MP mode)
|
||||
-----------------------------
|
||||
|
||||
Start the LMCache server with the blend engine:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
lmcache server --l1-size-gb 20 --eviction-policy LRU --engine-type blend
|
||||
|
||||
The ``blend`` engine composes a ``BlendModule`` into the server and requires
|
||||
``--supported-transfer-mode`` to be ``lmcache_driven`` or ``auto`` (the default). See
|
||||
:doc:`/mp/configuration` for the related server flags.
|
||||
|
||||
.. note::
|
||||
|
||||
The in-process CacheBlend documentation -- configuration knobs such as
|
||||
``LMCACHE_ENABLE_BLENDING`` and an end-to-end example -- is preserved in
|
||||
the Legacy section: :doc:`/kv_cache_optimizations/blending`.
|
||||
@@ -0,0 +1,32 @@
|
||||
.. _cachegen:
|
||||
|
||||
CacheGen
|
||||
===================
|
||||
|
||||
.. warning::
|
||||
|
||||
This page documents the behavior of LMCache's in-process mode (deprecated). Please consider using :doc:`LMCache MP mode </mp/index>` for better feature support and performance. For the MP mode equivalent of this page, see :doc:`/mp/cachegen`.
|
||||
|
||||
|
||||
Cachegen leverages KV cache's distributional properties to encode a KV cache into more compact bitstream representations with negligible decoding overhead.
|
||||
|
||||
|
||||
Configuring CacheGen in LMCache
|
||||
---------------------------------------
|
||||
|
||||
The settings should be very similar to :ref:`naive KV cache sharing <share_kv_cache>`.
|
||||
Only minor configurations need to be done to enable CacheGen.
|
||||
|
||||
To enable CacheGen in offline inference, we need to set:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Enable cachgen compression in LMCache
|
||||
os.environ["LMCACHE_REMOTE_SERDE"] = "cachegen"
|
||||
|
||||
To enable CacheGen in online inference, we need to set the ``remote_serde`` in the configuration yaml:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
# Enable cachgen compression in LMCache
|
||||
remote_serde: "cachegen"
|
||||
@@ -0,0 +1,13 @@
|
||||
Compression
|
||||
===========
|
||||
|
||||
KV cache compression can greatly reduces the size of the cache, which can be beneficial for both storage/memory usage and loading speed.
|
||||
Currently, we support the following compression algorithms:
|
||||
|
||||
- :ref:`CacheGen <cachegen>`: `CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving <https://dl.acm.org/doi/10.1145/3651890.3672274>`_
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
cachegen
|
||||
@@ -0,0 +1,11 @@
|
||||
KV Cache Optimizations
|
||||
======================
|
||||
|
||||
Techniques that improve KV cache reuse and efficiency beyond plain prefix
|
||||
caching.
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
||||
cacheblend
|
||||
segmented_prefill
|
||||
@@ -0,0 +1,116 @@
|
||||
Layerwise KV Transfer
|
||||
=====================
|
||||
|
||||
.. warning::
|
||||
|
||||
This page documents the behavior of LMCache's in-process mode (deprecated). Please consider using :doc:`LMCache MP mode </mp/index>` for better feature support and performance.
|
||||
|
||||
|
||||
The storage and loading of KV Cache on a layer granularity is a key optimization that allows for forward pass to "stagger" through its computation as each layer's KV Cache is received instead of only waiting to begin after the entire loading
|
||||
|
||||
CacheBlend is implemented on top of the layerwise codepath in order to pipeline recompute and loading to mask the latency of loading KV Cache.
|
||||
|
||||
.. image:: /_static/basic_codepath.svg
|
||||
:alt: Basic Codepath
|
||||
:class: scalable
|
||||
|
||||
.. raw:: html
|
||||
|
||||
<div style="text-align:center; margin:1em 0;">
|
||||
<a href="/_static/full_layerwise_diagram.svg" target="_blank">
|
||||
<img src="/_static/full_layerwise_diagram.svg"
|
||||
style="display:block; margin:auto; max-width:100%; height:auto;"/>
|
||||
</a>
|
||||
<div style="font-size:0.9em; color:#555; margin-top:0.5em;">
|
||||
Click to open full-size
|
||||
</div>
|
||||
</div>
|
||||
|
||||
Architecture Overview
|
||||
---------------------
|
||||
|
||||
**CacheEngine**
|
||||
The main orchestrator containing two primary generators:
|
||||
|
||||
* **Retrieval Generator** (N + 2 yields): Handles layer-by-layer KV cache loading with on-demand memory allocation
|
||||
* **Storage Generator** (N + 1 yields): Manages layer-by-layer KV cache saving with upfront CPU memory allocation
|
||||
|
||||
**LayerwiseGPUConnector**
|
||||
Manages GPU-CPU memory transfers with dedicated CUDA streams:
|
||||
|
||||
* **Load GPU Buffer**: Temporary GPU memory for CPU→GPU transfers (``use_gpu: true``)
|
||||
* **Store GPU Buffer**: Temporary GPU memory for GPU→CPU transfers (``use_gpu: true``)
|
||||
* **Nested Generators**: ``batched_to_gpu()`` and ``batched_from_gpu()`` handle actual memory operations
|
||||
|
||||
**StorageManager**
|
||||
Handles persistent storage operations:
|
||||
|
||||
* ``layerwise_batched_get()``: Asynchronous retrieval with ``.result()`` for request-level concurrency
|
||||
* ``batched_put()``: Stores memory objects to persistent backends
|
||||
|
||||
Execution Flow
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
The layerwise pipeline follows a numbered execution sequence:
|
||||
|
||||
**1. start_load_kv()**
|
||||
* Initializes Retrieval Generator via ``lmcache_engine.retrieve_layer()``
|
||||
* Performs setup (1st ``next()``) and loads layer 0 (2nd ``next()``)
|
||||
* Creates ``layerwise_retrievers`` list for ongoing layer processing
|
||||
|
||||
**2. wait_for_layer_load()** (repeated for each layer)
|
||||
* Advances Retrieval Generator via ``next()`` to process layer i
|
||||
* Triggers ``StorageManager.layerwise_batched_get()`` for async cache retrieval
|
||||
* Calls GPU Load Generator's ``batched_to_gpu()`` to transfer memory objects to GPU
|
||||
* **Last request in batch**: Synchronizes ``current_stream.wait_stream(load_stream)``
|
||||
|
||||
**3. save_kv_layer()** (repeated for each layer)
|
||||
* **First call only**: Creates Storage Generator with upfront CPU memory allocation
|
||||
* Advances Storage Generator via ``next()`` to process layer i
|
||||
* Calls GPU Store Generator's ``batched_from_gpu()`` to transfer GPU data to CPU
|
||||
* **First request in batch**: Synchronizes ``store_stream.wait_stream(current_stream)``
|
||||
|
||||
**4. wait_for_save()**
|
||||
* Finalizes Storage Generator with last ``next()`` call
|
||||
* Completes all ``StorageManager.batched_put()`` operations
|
||||
* Performs GPU Store Generator cleanup
|
||||
|
||||
Key Optimizations
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
**Pipelined Memory Operations**
|
||||
The system overlaps layer N+1 computation with layer N storage.
|
||||
|
||||
**Stream Synchronization**
|
||||
Three CUDA streams coordinate operations:
|
||||
|
||||
* ``current_stream``: vLLM's forward pass computation
|
||||
* ``load_stream``: KV cache loading operations
|
||||
* ``store_stream``: KV cache storing operations
|
||||
|
||||
**Batch-Level Coordination**
|
||||
Multiple requests are processed together with specialized synchronization:
|
||||
|
||||
* **First request**: Provides store stream synchronization to prevent GPU buffer corruption
|
||||
* **Last request**: Provides load stream synchronization to ensure KV cache availability
|
||||
|
||||
**Memory Allocation Strategies**
|
||||
* **Retrieval**: Layer-by-layer allocation
|
||||
* **Storage**: Upfront allocation for all layers
|
||||
|
||||
**Cache Key Management**
|
||||
Multi-layer cache engine keys use ``split_layers(N)`` to create per-layer kubernetes_deployment
|
||||
|
||||
Configuration
|
||||
~~~~~~~~~~~~~
|
||||
|
||||
Enable layerwise caching by setting:
|
||||
|
||||
.. code-block:: yaml
|
||||
|
||||
use_layerwise: true
|
||||
|
||||
The system automatically selects appropriate layerwise GPU connectors based on configuration:
|
||||
|
||||
* ``VLLMPagedMemLayerwiseGPUConnector``: For standard layerwise operations
|
||||
* ``VLLMBufferLayerwiseGPUConnector``: When blending is enabled
|
||||
@@ -0,0 +1,6 @@
|
||||
Segmented Prefill
|
||||
=================
|
||||
|
||||
.. note::
|
||||
|
||||
Documentation for segmented prefill is **coming soon**.
|
||||
Reference in New Issue
Block a user