逐层KV池

逐层KV池#

逐层模式是AscendStore KV池的一项优化，它逐层保存和加载KV缓存，而不是作为单个批量拷贝。通过将一层的传输与下一层的注意力计算流水线化，它减少了因必须等待整个KV缓存到达后才能进行任何前向推进而导致的停顿。

逐层模式在PD混合（kv_role: "kv_both"）和PD分离（kv_role: "kv_producer" / "kv_consumer"）场景下均有效。有关通用KV池架构和后端设置，请参阅KV池指南。

工作原理（简述）#

在没有逐层模式的情况下，请求的KV缓存会在完整前向传播完成后（或开始前）作为一次批量操作保存到池中（或从池中加载）。对于长提示词，这种批量传输会引入序列化停顿。

逐层模式在层粒度上拆分保存/加载：

  1. 保存（生产者 / kv_both）：计算完第i层的注意力后，该层的KV立即发送到池后端。下一层的计算与传输并行进行。

  2. 加载（消费者 / kv_both）：在计算第i层的注意力之前，系统等待第i层的KV从池中到达（wait_for_layer_load），然后继续。第i+1层的传输与第i层的注意力计算重叠。

最终效果：保存/加载延迟被分摊到整个前向传播过程中，而不是集中在一个阻塞步骤上。

前提条件#

逐层模式当前需要memcache后端（backend: "memcache"）。在继续之前，请安装并配置memcache_hybrid——有关memcache安装、配置文件（mmc-meta.conf / mmc-local.conf）和MetaService启动，请参阅KV池指南。

额外设置：

# Huge pages (required by memcache device transfer)
echo 200000 > /proc/sys/vm/nr_hugepages

# Source memcache environment
source /usr/local/memcache_hybrid/set_env.sh
source /usr/local/memfabric_hybrid/set_env.sh

# Uniform hashing across nodes
export PYTHONHASHSEED=0

配置#

将use_layerwise: true添加到AscendStoreConnector的额外配置中：

{
    "kv_connector": "AscendStoreConnector",
    "kv_role": "kv_both",
    "kv_connector_extra_config": {
        "backend": "memcache",
        "mooncake_rpc_port": "0",
        "use_layerwise": true
    }
}

对于PD分离，将"kv_role"更改为"kv_producer"或"kv_consumer"。

关键参数#

参数

默认值

描述

use_layerwise

false

启用逐层KV保存/加载。需要backend: "memcache"。

backend

"mooncake"

存储后端。逐层模式当前仅支持"memcache"。

mooncake_rpc_port

"0"

调度器↔工作节点查找服务的RPC端口。使用"0"自动分配，或为每个实例指定唯一端口。

layerwise_prefetch_layers

1

在计算前沿之前预取的层数。更高的值可以提高重叠度，但会消耗更多内存。

layerwise_max_transfer_blocks

0（无限制）

每次传输批次的KV块最大数量。

layerwise_max_transfer_bytes

0（无限制）

每次传输批次的最大字节数。

h2d_stagger_us

0

跨TP rank的H2D拷贝之间的交错延迟（微秒），以避免总线争用。

discard_partial_chunks

true（非逐层）/ false（逐层）

是否丢弃不完整块边界的KV。逐层模式默认为false以保留部分层。

使用场景#

PD混合（kv_both）#

单个vLLM实例同时充当生产者和消费者。池作为共享前缀缓存：已完成请求的KV逐层保存，具有重叠前缀的新请求逐层加载KV。无需代理。

export ASCEND_RT_VISIBLE_DEVICES=0,1

python -m vllm.entrypoints.openai.api_server \
    --model /path/to/DeepSeek-V2-Lite \
    --port 8100 \
    --trust-remote-code \
    --enforce-eager \
    --no-enable-prefix-caching \
    --tensor-parallel-size 1 \
    --max-model-len 4096 \
    --max-num-batched-tokens 4096 \
    --kv-transfer-config '{
        "kv_connector": "AscendStoreConnector",
        "kv_role": "kv_both",
        "kv_connector_extra_config": {
            "backend": "memcache",
            "mooncake_rpc_port": "0",
            "use_layerwise": true
        }
    }'

直接向端口8100发送请求——无需代理。

PD分离（kv_producer + kv_consumer）#

独立的预填充器和解码器实例。预填充器逐层保存KV；解码器逐层加载KV。逐层代理通过其/v1/metaserver端点协调请求路由和逐层KV放置。

预填充器：

export ASCEND_RT_VISIBLE_DEVICES=0,1

python -m vllm.entrypoints.openai.api_server \
    --model /path/to/DeepSeek-V2-Lite \
    --port 8100 \
    --trust-remote-code \
    --enforce-eager \
    --tensor-parallel-size 1 \
    --max-model-len 4096 \
    --kv-transfer-config '{
        "kv_connector": "AscendStoreConnector",
        "kv_role": "kv_producer",
        "kv_connector_extra_config": {
            "backend": "memcache",
            "mooncake_rpc_port": "0",
            "use_layerwise": true
        }
    }'

解码器：

export ASCEND_RT_VISIBLE_DEVICES=2,3

python -m vllm.entrypoints.openai.api_server \
    --model /path/to/DeepSeek-V2-Lite \
    --port 8200 \
    --trust-remote-code \
    --enforce-eager \
    --tensor-parallel-size 1 \
    --max-model-len 4096 \
    --kv-transfer-config '{
        "kv_connector": "AscendStoreConnector",
        "kv_role": "kv_consumer",
        "kv_connector_extra_config": {
            "backend": "memcache",
            "mooncake_rpc_port": "0",
            "use_layerwise": true
        }
    }'

逐层代理（与标准分离代理不同——提供/v1/metaserver服务）：

python examples/disaggregated_prefill_v1/load_balance_proxy_layerwise_server_example.py \
    --host 127.0.0.1 \
    --port 9000 \
    --prefiller-hosts 127.0.0.1 \
    --prefiller-ports 8100 \
    --decoder-hosts 127.0.0.1 \
    --decoder-ports 8200

注意： 代理--host不能是0.0.0.0（通配符）。解码器会连接回host:port/v1/metaserver，因此请使用可访问的IP。

向代理发送请求：

curl -s http://127.0.0.1:9000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "/path/to/DeepSeek-V2-Lite",
        "prompt": "Hello, my name is",
        "max_tokens": 32,
        "temperature": 0.0
    }'

调优#

预取深度#

增加layerwise_prefetch_layers（默认1）以在计算前沿之前预取更多层。这增加了传输/计算重叠，但会使用更多临时缓冲区。典型值：1–4。

传输批处理#

使用layerwise_max_transfer_blocks或layerwise_max_transfer_bytes来限制每个传输批次的大小。这可以防止单个大层独占传输总线。设置为0（默认）表示无限制。

H2D交错#

在多TP部署中，所有TP rank的H2D（主机到设备）拷贝可能会在PCIe/HCCS总线上争用。设置h2d_stagger_us以分散它们（例如，100表示rank之间100微秒的交错）。

支持的模型#

逐层模式与 MLA（mla_v1）和 SFA（sfa_v1）注意力后端集成。支持 DeepSeek-V2/V3 及其他基于 MLA 的模型。

基础全注意力（attention_v1）及所有上下文并行（CP）变体（mla_cp、sfa_cp、attention_cp）尚未集成逐层等待/保存调用。逐层+CP 是未来工作。

限制#

  • 后端：逐层模式仅支持 memcache（mooncake 和 yuanrong 不支持 use_layerwise）。

  • 混合 KV 缓存：不支持——当模型具有多个 KV 缓存组族（混合 MLA + 滑动窗口注意力）时，逐层模式会抛出 NotImplementedError。

  • 上下文并行：逐层模式尚未与 CP 注意力后端集成。

  • PD 分离代理：使用 kv_producer / kv_consumer 时，需要专用的逐层代理（load_balance_proxy_layerwise_server_example.py）——标准分离代理不提供 /v1/metaserver 端点。