# [Feature] 移植 EXP KV 压缩方案到 v0.15.1-dev（v1 引擎）

**Issue ID**: #001
**Status**: Open
**Priority**: High
**Created**: 2026-02-23
**Updated**: 2026-02-24
**Assignee**: Unassigned
**Labels**: feature, enhancement, kv-compression, v1, v0.15.1, exp

---

## 描述

将 EXP 分支（`vllm-exp` 的 `v0.9.2-kvpressv2`）中已实现的 v1 KV 压缩（token-shared Top‑K + compaction，含 chunked-prefill prompt-end one-shot）移植并接入到 `vllm` 的 `v0.15.1-dev` 分支。

目标是：在 **不破坏 v0.15.1 现有调度/缓存/connector 语义** 的前提下，引入可控开关的 KV 压缩能力，先跑通最小闭环，再逐步放开组合与优化。

## 需求背景

EXP 分支 KV 压缩已在单个大 commit 中完成接入（参考 `vllm-exp@332965b74`），但 `v0.15.1-dev` 的 v1 代码结构、调度输出、KVCacheManager 以及 attention backend 的接口拆分均有变化，不能直接搬目录；需要按“端到端链路”重新对齐接入点。

## 用户故事

作为 vLLM v1 引擎的开发/维护者，我希望在 `v0.15.1-dev` 上也能启用 KV 压缩，以便在长 prompt / 大批量 prefill 时降低 KV 占用并提升吞吐，同时保持质量与稳定性可控。

## 功能需求

### 基本功能

1. **Feature flag + 参数**
   - `VLLM_ENABLE_KV_COMPRESSION=1` 可启用 KV 压缩（默认关闭）。
   - 支持 `topk` policy 相关参数（prompt ratio/budget、protected prefix/suffix、SnapKV window 等），与 EXP 分支保持一致。

2. **端到端正确性（核心）**
   - 引入 `request.num_kv_tokens`，并保证：
     - 逻辑进度：`num_computed_tokens`（token/RoPE 位置）
     - KV 实际长度：`num_kv_tokens`（KV cache 中保留条目数）
   - KV 写入位置与逻辑位置解耦：slot_mapping 需基于 `kv_positions` 而不是逻辑 `positions`。
   - KV block 分配需在压缩开启时按 `num_kv_tokens`（并兼容 connector/external tokens 路径）。

3. **两套压缩路径**
   - **Scheme 3（chunked-prefill）**：prompt ingest 完成时计算一次 keep indices，decode 前一次性 compaction。
   - **Scheme 1/2（非 chunked）**：每 step 内做 token-shared 选择 + reshape_and_cache 重写（可先在最小闭环后再放开）。

4. **平台/组合约束 gate**
   - 先限制到 CUDA/ROCm + FlashAttention v1 backend。
   - 对不兼容组合进行显式禁用/报错（prefix caching / sliding window / full cuda graph / spec decode / context parallel / KV sharing 等），避免隐式错位。

### 高级功能（可选，后续迭代）

5. **Prompt compaction 后释放 tail blocks（省显存）**
   - 参考 EXP：在 KV compaction 后尝试 truncate/free 不再需要的尾部 blocks（需要在 v0.15.1 的 single_type manager / coordinator 增补 hook）。

6. **逐步放开组合**
   - 在正确性跑通后，逐步评估与 DCP/PCP、prefix caching、connector、spec decode 等组合的可行性与改造成本。

## 技术方案

### 1. 参考实现（EXP）

单 commit：`vllm-exp@332965b74 feat(v1): 集成 KV 压缩完整实现`

该提交主要改动点（用于对照移植）：
- `vllm-exp/vllm/envs.py`：新增 KV compression env
- `vllm-exp/vllm/v1/request.py`：新增 `num_kv_tokens`
- `vllm-exp/vllm/v1/core/sched/output.py`：透传 `num_kv_tokens`
- `vllm-exp/vllm/v1/core/sched/scheduler.py`：gate + 记账更新 `num_kv_tokens`
- `vllm-exp/vllm/v1/core/kv_cache_manager.py`：allocate_slots 按 `num_kv_tokens`
- `vllm-exp/vllm/v1/worker/gpu_input_batch.py`：维护 `num_kv_tokens_cpu`
- `vllm-exp/vllm/v1/worker/gpu_model_runner.py`：计算 `kv_positions` + slot_mapping/seq_lens
- `vllm-exp/vllm/v1/attention/backends/flash_attn.py`：metadata 扩展 + hooks
- `vllm-exp/vllm/v1/kv_compression/*`：Top‑K / SnapKV / Triton gather+compact / prompt-end

### 2. v0.15.1 接入要点（对照清单）

#### A) 状态与透传
- `vllm/vllm/v1/request.py`：补 `num_kv_tokens`
- `vllm/vllm/v1/core/sched/output.py`：New/CachedRequestData 补 `num_kv_tokens`
- `vllm/vllm/v1/worker/gpu_input_batch.py`：CachedRequestState/InputBatch 补 `num_kv_tokens(_cpu)`

#### B) Scheduler 记账 + KV 分配
- `vllm/vllm/v1/core/sched/scheduler.py`：
  - 在 `_update_after_schedule()` 中按 step 推进 `num_kv_tokens`（需记录 `start_pos`）
  - 处理 cache-hit/connector/抢占的初始化与重置
- `vllm/vllm/v1/core/kv_cache_manager.py`：
  - allocate_slots 在压缩开启时按 `num_kv_tokens` 计算 need_slot（兼容 external computed tokens）

#### C) Worker：slot_mapping 与 seq_lens
- `vllm/vllm/v1/worker/gpu_model_runner.py`：
  - `_prepare_inputs()`：保留逻辑 `positions_np`，新增物理 `kv_positions_np`
  - `block_table.compute_slot_mapping()`：开启压缩时使用 `kv_positions_np`
  - `seq_lens`：开启压缩时用 `num_kv_tokens + scheduled_len`
  - **注意**：`discard_request_mask` 仍应按逻辑长度判断（v0.15.1 里 `seq_lens` 被复用，需避免误用 KV 长度导致“错误丢 token”）

#### D) Attention：metadata + hooks
- `vllm/vllm/v1/attention/backends/flash_attn.py`：
  - 扩展 `FlashAttentionMetadata`（must_keep/topk_budget/prompt_end/...）
  - builder 填充 KV compression metadata
  - 在 `forward()` 中接入 compaction / prompt-end payload hooks

#### E) KV compression 模块
- 新增/移植目录：`vllm/vllm/v1/kv_compression/`（从 EXP 同名目录移植并对齐 import）

## 实现计划

### Phase 1: Plumbing（端到端链路跑通）
- [x] envs + scheduler gate（先限制组合）
- [x] 引入 `num_kv_tokens`（request/output/worker 透传）
- [x] scheduler 记账推进 `num_kv_tokens`
- [x] kv_cache_manager allocate_slots 按 `num_kv_tokens`
- [x] worker slot_mapping/seq_lens 切换到 `kv_positions`
- [x] 修复 v0.15.1 的 `discard_request_mask` 逻辑长度判断问题

### Phase 2: Scheme 3（chunked-prefill prompt-end one-shot）
- [x] prompt-end payload（gather K + SnapKV scoring + Top‑K indices）
- [x] runner stash + decode 前一次性 compaction
- [ ] 基础正确性验证（prefill→decode 质量不崩）

### Phase 3: Scheme 1/2（非 chunked 每 step compaction）
- [x] step 内 must_keep/topk_budget + dst_slots
- [x] reshape_and_cache 重写（仅 rewrite moved tokens）
- [ ] 性能回归与参数调优

### Phase 4: 优化与扩展（可选）
- [ ] compaction 后释放 tail blocks（truncate hook）
- [ ] 逐步评估/放开更多组合（DCP/PCP、connector、prefix caching…）

## 相关文件

- `vllm-exp/vllm/v1/kv_compression/`（参考实现）
- `vllm/vllm/v1/core/sched/scheduler.py`
- `vllm/vllm/v1/core/sched/output.py`
- `vllm/vllm/v1/core/kv_cache_manager.py`
- `vllm/vllm/v1/request.py`
- `vllm/vllm/v1/worker/gpu_model_runner.py`
- `vllm/vllm/v1/worker/gpu_input_batch.py`
- `vllm/vllm/v1/attention/backends/flash_attn.py`
- `vllm/vllm/envs.py`

## 进展记录

### 2026-02-23 22:32
- [x] 完成需求分析与改动点梳理
- [x] 完成 Phase 1（plumbing，代码接入）
- [ ] 完成 Scheme 3（chunked-prefill，待端到端验证）
- [ ] 完成 Scheme 1/2（per-step，待端到端验证/调优）

### 2026-02-23 23:43
- 移植难度评估：
  - 最小闭环（Phase 1 + Scheme 3 + 严格 gate）：中等偏难（需要多处状态/语义对齐，但可控）。
  - 完整对齐 EXP（含 Scheme 1/2 + tail-block truncate/free + 放开更多组合）：难（需要更深的调度/缓存/attention 语义改造与回归）。
- 关键风险点（建议在 Phase 1 先对齐/加断言）：
  - `seq_lens` 的复用：开启压缩后 `seq_lens`=KV 长度，但 `discard_request_mask` 等逻辑仍应基于**逻辑 token 长度**（`num_computed_tokens`/positions），否则可能 silent drop token。
  - `num_kv_tokens` 初始化/重置：cache-hit、connector/external tokens、抢占/恢复时必须一致，否则 allocate_slots/slot_mapping 易错位。
  - slot_mapping：开启压缩时必须基于 `kv_positions`（物理 KV 位置），不能继续用逻辑 `positions`。
  - KV 分配：allocate_slots 需按 `num_kv_tokens` 计数，并兼容 external computed tokens 的“已占用 KV”语义。
- 推荐落地顺序：
  1) Phase 1 plumbing + 严格 gate + correctness asserts（先保证开关切换不破坏现有路径）。
  2) Scheme 3（prompt-end one-shot）先跑通正确性与稳定性。
  3) Scheme 1/2（per-step）与 tail-block truncate/free 放到后面做性能/显存优化与组合扩展。

### 2026-02-24 01:27
- 已落地 v0.15.1-dev plumbing + gate：
  - `envs` 增加 KV compression 开关与参数
  - `Request/SchedulerOutput/InputBatch/GPUModelRunner` 贯通 `num_kv_tokens`
  - Scheduler 侧接入 `scheduler_accounting.update_num_kv_tokens_after_schedule`
  - `KVCacheManager.allocate_slots` 在开启压缩时按 `num_kv_tokens` 分配
  - Worker 侧引入 `kv_positions`，slot_mapping 基于物理 KV 位置而不是逻辑 positions
  - Attention backend runtime gate：KV compression 仅允许 `FLASH_ATTN`
- 已接入 FlashAttention metadata + hooks：
  - `CommonAttentionMetadata/FlashAttentionMetadata` 扩展 KV compression 字段
  - FlashAttention forward 内接入 prompt-end payload（scheme 3）与 per-step compaction（scheme 1/2）hooks
  - prompt-end payload -> runner stash -> decode 前 one-shot prompt compaction（scheme 3）
- 待验证项：
  - 实际推理端到端验证（prefill→decode）与质量回归
  - 与更多组合（prefix caching / sliding window / spec decode / fullgraph / CP）仍保持禁用

### 2026-02-24 14:10
- 修复移植差异导致的关键问题（review 阶段发现）：
  - `chunked_prefill_enabled` 字段名在 v0.15.1 为 `enable_chunked_prefill`（否则 scheme 3 永远不触发或直接报错）
  - `KVCacheManager.allocate_slots` 需基于 `num_kv_tokens` 分配，否则压缩后仍按逻辑长度分配、收益被吃掉
  - ROCm 下 flash-attn KV cache 的 block_size 维度不同，需修正 `cache_block_size` 取值

### 2026-02-24 14:25
- 开始按“逐个 commit”方式推进（便于 review/回溯）：
  - Commit 1：`feat(kvpress): 新增 KV 压缩环境变量开关与参数`（`cafabeeb5`，`vllm/envs.py`，已提交）
  - Commit 2：`feat(kvpress): 调度输出透传 num_kv_tokens`（`c44fcded3`，`vllm/v1/request.py` + `vllm/v1/core/sched/output.py`，已提交）
  - Commit 3：`feat(kvpress): 扩展 InputBatch 请求状态`（`d41ca1284`，`vllm/v1/worker/gpu_input_batch.py`，已提交）
  - Commit 4：`feat(kvpress): 新增 Top-K budget 与选择工具`（`3da2c8293`，`vllm/v1/kv_compression/{budget,slot_mapping,topk_select}.py`，已提交）
  - Commit 5：`feat(kvpress): 新增 SnapKV 打分与 Triton KV 算子`（`87b788bd9`，`vllm/v1/kv_compression/{snapkv_*,kv_cache_*}.py`，已提交）
  - Commit 6：`feat(kvpress): 增加调度侧 KV 长度记账`（`b0911b244`，`vllm/v1/core/sched/scheduler.py` + `vllm/v1/kv_compression/scheduler_accounting.py`，已提交）
  - Commit 7：`feat(kvpress): KVCacheManager 按 num_kv_tokens 分配 slots`（`3d4f87530`，`vllm/v1/core/kv_cache_manager.py`，已提交）
  - Commit 8：`feat(kvpress): Runner 接入 KV 位置与注意力元数据`（`dbcb03763`，`vllm/v1/worker/gpu_model_runner.py` + runner 辅助模块，已提交）
  - Commit 9：`feat(kvpress): FlashAttention 接入 KV 压缩 hooks`（`da35d84f3`，`vllm/v1/attention/backends/flash_attn.py` + hooks 模块，已提交）

## 验收标准

- [ ] `VLLM_ENABLE_KV_COMPRESSION=0` 不影响现有行为
- [ ] `VLLM_ENABLE_KV_COMPRESSION=1` 在支持平台/组合下可稳定跑通 prefill→decode
- [ ] slot_mapping/seq_lens 不错位（无越界、无错写、无随机崩溃）
- [ ] 关键互斥组合有明确 gate（避免 silent wrong）
- [ ] 基础回归测试通过（至少覆盖 v1 引擎常用路径）

---

## 备注

- 参考 commit：`vllm-exp@332965b74`
- 建议先以 “FlashAttention + 单机单卡/无 CP-DCP + 禁用 prefix caching/sliding window/spec decode/fullgraph” 跑通最小闭环，再逐步放开。
