docs(issue): 新增 #001 - KV 压缩移植跟踪

.issues/open/001-feature-kv-compression-port-to-v0.15.1.md

+# [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” 跑通最小闭环，再逐步放开。