kvpress_v2.md

# KVPress v2 PR 整理
# 容器位置 注意实际跑的容器的位置是在/public/home/lixh6/laibao/ssh/kvpress.sh 脚本对应的容器
## PR1：feat: kvpress新增 KV 压缩配置开关（默认关闭）
**目的**：提供 KV compression 的配置入口，默认关闭，不改变现有行为。  
**范围**：`vllm/envs.py`  
**内容**：新增 KV compression 的开关与参数（policy、ratio/budget、protected 前后缀、SnapKV window 等）。  
**新增环境变量说明**：
- `VLLM_ENABLE_KV_COMPRESSION`：总开关，开启 v1 token-shared KV 压缩。
- `VLLM_KV_COMPRESSION_POLICY`：压缩策略，目前仅支持 `topk`。
- `VLLM_KV_COMPRESSION_PROMPT_RATIO`：按比例保留 prompt token（只对非保护区生效）。
- `VLLM_KV_COMPRESSION_PROMPT_BUDGET`：按数量保留 prompt token，`>=0` 时优先生效，覆盖 ratio。
- `VLLM_KV_COMPRESSION_PROTECTED_PREFIX`：强制保留 prompt 前缀 N 个 token。
- `VLLM_KV_COMPRESSION_PROTECTED_SUFFIX`：强制保留 prompt 后缀 N 个 token。
- `VLLM_KV_COMPRESSION_KEEP_LAST_TOKEN`：强制保留 prompt 最后一个 token。
- `VLLM_KV_COMPRESSION_SNAPKV_WINDOW`：SnapKV-like scoring 的窗口大小。
- `VLLM_KV_COMPRESSION_SNAPKV_USE_TRITON_ROCM`：ROCm 是否用 Triton 版本 scoring（否则走 torch 参考实现）。
- `VLLM_KV_COMPRESSION_TOPK_PER_LAYER`：是否每层各自做 Top-K（默认跨层共享一次选择）。
- `VLLM_KV_COMPRESSION_ASYNC_WRITEBACK`：compaction 写回是否用独立 CUDA stream（异步）。
- `VLLM_KV_COMPRESSION_FREE_TAIL_BLOCKS`：压缩后是否释放多余尾部 KV blocks（影响并发上限）。
**风险**：低（仅新增配置项，无功能路径变化）。

## func VLLM_KV_COMPRESSION_FREE_TAIL_BLOCKS
**定位**：KV compression 的可选“资源回收”开关。KV compression 会让 `num_kv_tokens`（KV 实际长度）小于逻辑长度，但 request 之前申请到的 KV blocks 不会自动缩回；该开关用于在合适时机把“压缩后不再需要的尾部 blocks”归还给 block pool，提高并发上限/降低被抢占概率。  

**它解决的具体问题**：
- **不回收**：request 在 prompt 很长但压缩很激进、且 decode 很短的场景，会长期占着一段“不会再被读取/写入”的尾部 blocks，导致全局可用 blocks 变少（并发下降）。
- **回收后**：尾部 blocks 立刻回到 pool，其他 request 可以在同一个 engine step 内复用这些 blocks（写入会覆写旧 KV）。  

**触发位置（Scheduler）**：`vllm/v1/core/sched/scheduler.py`  
- 条件：`self.kv_compression_enabled` + `VLLM_KV_COMPRESSION_FREE_TAIL_BLOCKS=1` + prompt 已结束（`request.num_computed_tokens == request.num_prompt_tokens`）。  
- 动作：调用 `self.kv_cache_manager.truncate_to_num_tokens(request_id, request.num_kv_tokens)` 尝试截断该 request 的 blocks（best-effort；返回 True 表示确实释放了 blocks）。  
- 若释放成功：把 `request_id` 加入 `force_replace_block_ids`，在本 step 的 `SchedulerOutput` 中强制 worker **replace** 该请求的 `block_ids`（而不是 append）。  

**释放链路（best-effort，且是“即时释放”）**：
- `vllm/v1/core/kv_cache_manager.py:truncate_to_num_tokens`  
  → `vllm/v1/core/kv_cache_coordinator.py:truncate_to_num_tokens`  
  → `vllm/v1/core/single_type_kv_cache_manager.py:FullAttentionManager.truncate_to_num_tokens`  
  → `vllm/v1/core/block_pool.py:free_blocks`（blocks 立刻回到 free queue，可被本 step 其它请求复用）。  

**为什么必须 `force_replace_block_ids`（否则会有 correctness 风险）**：
- 截断 blocks 会让 `block_ids` 变短；如果 worker 仍按“增量 append”更新 block table，旧的尾部 block id 仍会残留在 worker 的 `block_table` 中。  
- 由于这些 blocks 已经回到 pool，可能马上被其它请求重新分配并写入（覆写旧 KV）；残留引用会导致后续 attention **读到错误的 KV**。  
- 因此 scheduler 复用 `resumed_from_preemption` 的语义来实现“整行替换”：  
  - `vllm/v1/core/sched/output.py`：`CachedRequestData.resumed_from_preemption=True` 表示用 `new_block_ids` **replace**（不是 append）。  
  - `vllm/v1/core/sched/scheduler.py:_make_cached_request_data`：当 `req_id ∈ force_replace_block_ids` 时置 True。  
  - `vllm/v1/worker/gpu_model_runner.py`：`resumed_from_preemption` 分支里 `req_state.block_ids = new_block_ids`，并用 `block_table.add_row(...)` 替换整行。  

**关闭该开关时的行为（通常不影响精度，但影响资源）**：
- 不会主动 free tail blocks；这些 blocks 会一直被该 request 占用直到请求结束。  
- 精度通常不变：因为 KV compression 下 attention 的有效读写由 `num_kv_tokens`/`kv_positions` 控制，尾部 blocks 不会再参与计算；但它会降低全局可用 blocks，从而影响并发/抢占。  
- 另外：如果 decode 很长，`num_kv_tokens` 会持续增长，之前“多占着”的 blocks 可能会被后续 decode 写入覆盖，浪费会变小；因此该开关的收益更多体现在“高压缩 + 短输出 + 多并发”的场景。  

**与 chunked prefill（scheme 3）的关系 / 风险提示**：
- scheme 3 的 “one-shot prompt compaction” 是在 **第一次 decode 前**由 runner 执行（`vllm/v1/worker/gpu_model_runner.py:_maybe_apply_kv_compression_prompt_compaction`）；在此之前，prompt KV 仍按原始连续布局存放。  
- 如果在 compaction 之前就 truncate blocks，会误释放仍需读取/搬运的 prompt KV（可能直接错答/崩溃）。  
- 因此：chunked prefill 场景建议先保持 `VLLM_KV_COMPRESSION_FREE_TAIL_BLOCKS=0`；或引入“compaction 已完成”的显式信号/状态位后再允许 truncate。  

## PR2：feat: kvpress新增 KV 压缩预算计算模块
**目的**：提供 KV compression 预算计算的公共逻辑，供调度与 runner 复用。  
**范围**：`vllm/v1/kv_compression/__init__.py`、`vllm/v1/kv_compression/budget.py`  
**内容**：新增 prompt 保留预算、must-keep 统计、step 级 Top-K 预算等基础工具函数。  
**核心函数说明**：
- `count_prompt_must_keep_in_range(...)`：统计区间内必须保留的 prompt token（protected 前后缀 + 可选 last token）。
- `compute_topk_budget_step(...)`：计算单步 Top-K 预算，按候选 token 前缀比例分配。
- `compute_prompt_topk_keep_total(...)`：计算整个 prompt 的候选 token 总保留数。
- `compute_prompt_keep_len(...)`：总保留长度 = must-keep + Top-K 保留，并 clamp 到 [0, prompt_len]。
**要点**：预算以 `prompt_budget>=0` 优先，其次按 ratio；ratio 会 clamp 到 [0,1]，并用四舍五入分配。  
**风险**：低（纯计算工具，无行为变更）。

## PR3：feat: kvpress新增 SnapKV 打分与 KV compaction Triton 内核
**目的**：提供 KV compression 的核心 Triton 算子（打分 + KV compaction）。  
**范围**：`vllm/v1/attention/kv_compression/kv_cache_triton.py`、`vllm/v1/attention/kv_compression/snapkv_triton.py`  
**内容**：新增 SnapKV-like 打分内核与 KV cache gather/前移压缩内核。  
**不开 chunked prefill 的路径**：
- 在每个 step 内完成 Top-K 选择 + `reshape_and_cache` 写回压缩后的 KV（即时压缩）。
- 该路径使用 `snapkv_triton.py` 进行打分，不使用 `kv_cache_triton.py` 的 gather/前移。
**开启 chunked prefill（scheme 3）的路径**：
- 在最后一个 prompt chunk 仅计算并缓存全局 Top-K 索引（通过 `gather_k_to_packed_triton`）。
- 在第一次 decode 前执行一次性 KV 前移压缩（`front_compact_inplace_fa_triton`）。
**风险**：中（涉及 Triton 内核与 KV 写回路径）。
**待处理问题**：
- `vllm/v1/attention/kv_compression/snapkv_triton.py:285`：`k_eff_end <= k_beg` 直接 return，`PROTECT_LAST` 在后面，保护逻辑会被跳过；当前 v1 调用显式 `protect_last=False`，但仍建议修复以防后续开启。

## PR4：feat: kvpress新增 KV 压缩状态与元数据打通
**目的**：打通 KV 压缩在 request / scheduler output / input batch 的状态承载。  
**范围**：`vllm/v1/request.py`、`vllm/v1/core/sched/output.py`、`vllm/v1/worker/gpu_input_batch.py`、`vllm/v1/worker/block_table.py`  
**内容**：增加 `num_kv_tokens` 与 prompt compaction 元数据（idx/keep_len 等）的流转与缓存。  
**风险**：低（状态字段新增，尚未触及核心执行逻辑）。

## PR5：feat: kvpress新增 KV cache 申请/截断支持
**目的**：让 KV cache 的容量与 `num_kv_tokens` 对齐，并提供压缩后释放尾部 block 的能力。  
**范围**：`vllm/v1/core/kv_cache_manager.py`、`vllm/v1/core/kv_cache_coordinator.py`、`vllm/v1/core/single_type_kv_cache_manager.py`  
**内容**：
- `KVCacheManager` 在开启 KV compression 时按 `num_kv_tokens` 计算需要的 slots，避免按逻辑长度超分配。
- `KVCacheCoordinator`/`SingleTypeKVCacheManager` 增加 `truncate_to_num_tokens` 接口，用于压缩后回收尾部 blocks。
- `FullAttentionManager` 实现尾部 block 的实际释放逻辑。
**风险**：中（涉及 cache 分配/释放路径，需结合调度器调用点验证）。

## PR6：feat: kvpress新增调度层 KV 压缩逻辑
**目的**：在调度器侧实现 KV 压缩的核心策略与状态更新。  
**范围**：`vllm/v1/core/sched/scheduler.py`  
**内容**：
- 增加 KV compression 的开关与不兼容项校验（TPU、sliding window、prefix cache、CUDA graph、spec decode）。
- 维护 `num_kv_tokens` 的更新规则（区分 chunked prefill 与非 chunked prefill）。
- 压缩后释放尾部 block 时，强制替换 block IDs（避免 worker 端 append 导致残留）。
- 预抢占时同步清零 `num_kv_tokens`，保证状态一致性。
**核心逻辑说明**：
- 这部分代码不做真实 KV compaction，而是更新“实际 KV 长度”账本 `num_kv_tokens`。
- 未开启压缩时：`num_kv_tokens` 与 `num_computed_tokens` 同步，行为与原 vLLM 一致。
- 开启压缩时：
  - decode token 全保留；
  - prompt token 由 “must-keep(前缀/后缀/last)” + “Top-K 预算(ratio/budget)” 两部分组成；
  - chunked prefill 下，prompt 未结束时不截断，最后一个 chunk 结束时一次性计算保留长度。
**风险**：中（调度核心逻辑变更，需结合后续 runner/attention 路径验证）。
**待关注问题**：
- KV compression 开启 + KVConnector 同步命中（`load_kv_async=False`）时，`num_kv_tokens` 可能未被正确初始化，导致 KV 写入 offset 错位、覆盖已加载 KV 的风险。可在 RUNNING 初始化时将 `num_kv_tokens` 对齐 `num_computed_tokens`，或明确禁止该路径。
- 开启 chunked prefill（scheme 3）时，如果 `VLLM_KV_COMPRESSION_FREE_TAIL_BLOCKS=1`，调度器可能在“first decode step 调度”阶段提前 truncate blocks（此时 KV 尚未做一次性前移压缩），存在误释放仍需读取的 prompt KV 的风险；建议先保持该开关关闭，或增加“compaction 已完成”信号后再允许 truncate。

## PR7：feat: kvpress新增 runner 侧 KV 压缩状态/位置打通
**目的**：让 GPU runner 能识别 `num_kv_tokens` 并为压缩后的 KV 写入位置提供基础支撑。  
**范围**：`vllm/v1/worker/gpu_model_runner.py`  
**内容**：
- 新增 `kv_positions_cpu/kv_positions_np` 缓冲区，用于区分“逻辑位置”和“KV 写入位置”。
- 新请求与缓存请求同步 `num_kv_tokens` 到 runner 与 input_batch。
- 结合 `resumed_from_preemption`，在需要时用 `add_row` 替换 block_ids。
**风险**：中（runner 关键路径变更，需要与后续 slot mapping / attention 压缩逻辑联动验证）。

## PR8：feat: kvpress runner 侧按 num_kv_tokens 计算 KV 写入位置
**目的**：让 runner 能区分“逻辑 token 位置”和“KV cache 实际写入位置”，为后续 Top‑K 压缩/重排打基础。  
**范围**：`vllm/v1/worker/gpu_model_runner.py`  
**内容**：
- 在 KV compression 开启时计算 `kv_positions = num_kv_tokens + step_arange`，作为本 step 的 KV 写入位置。
- slot mapping 从 `positions` 切换为 `kv_positions`（仅在开启 KV compression 时）。
- `seq_lens` 从逻辑长度切换为 KV 长度（`num_kv_tokens`），保证 attention 看到的是“实际 KV 长度”。
- 开启 KV compression 时禁用 cascade attention（common-prefix 假设不再成立）。
- runner 侧 fail-fast 校验：不支持 full CUDA graph；仅支持 `FLASH_ATTN_VLLM_V1`；暂不支持 Mamba。
**风险**：低~中（默认关闭不影响现有行为；开启后影响 KV 写入/attention 长度计算，并对不兼容配置直接报错）。
**待处理问题**：
- pooling 请求在开启 KV compression 时可能卡死：当前 `_pool` 使用 `seq_lens == prompt_len` 判断是否输出 embedding，但 KV compression 下 `seq_lens` 代表 KV cache 长度（`num_kv_tokens`），可能永远小于 `prompt_len`。建议 pooling 完成条件改用逻辑长度（`num_computed_tokens`），或为 pooling 单独维护一份 logical `seq_lens`。

## PR9：feat: kvpress runner 侧生成 Top-K 压缩元数据
**目的**：在 runner 侧为 “topk” KV 压缩策略准备必要的元数据（must-keep 掩码、Top-K 预算、chunked prefill prompt-end 标记），并按需拷到 GPU。  
**范围**：`vllm/v1/worker/gpu_model_runner.py`  
**内容**：
- 新增 KV compression 相关元数据 buffer（CPU pinned + GPU）：per-token `must_keep`、per-request `topk_budget`，以及 chunked prefill 的 `prompt_end/prompt_len/topk_keep`。
- 非 chunked prefill：在 `_prepare_inputs` 计算本 step 每个 token 的 `must_keep`（decode 全保留 + prompt 保护区 + 可选 last token），并计算每个请求的 `topk_budget`；同时维护 `kv_compression_needs_compaction`（fast path：decode-only 且预算为 0 时跳过后续 score/topk/compaction）。
- chunked prefill（scheme 3）：在最后一个 prompt chunk 标记 `prompt_end`，并计算全 prompt 的 Top-K keep 数（仅生成元数据；一次性 compaction 在后续 PR 补齐）。
- 将上述元数据按需 `non_blocking` 拷贝到 GPU（仅在需要 compaction 或 chunked prefill prompt-end 时）。
**元数据 buffer 含义**（`*_cpu` 为 CPU pinned tensor，`*_np` 为其 numpy view，未带后缀者为 GPU tensor）：
- `kv_compression_must_keep_*`：per-token `bool[max_num_tokens]`，本 step 展平后的每个 token 是否“必须保留 KV”（decode 永远保留；prompt 的 protected prefix/suffix/last token 必保留）。
- `kv_compression_topk_budget_*`：per-request `int32[max_num_reqs]`，本 step 中每个请求的 Top‑K 预算（除 must-keep 外，还允许从候选 prompt token 里额外保留多少个）。
- `kv_compression_prompt_end_*`：per-request `bool[max_num_reqs]`，仅用于 chunked prefill：该 step 是否结束 prompt（跨过 prompt 末尾），用于触发后续“一次性 compaction”准备。
- `kv_compression_prompt_lens_*`：per-request `int32[max_num_reqs]`，仅用于 chunked prefill：prompt 总长度。
- `kv_compression_prompt_topk_keep_*`：per-request `int32[max_num_reqs]`，仅用于 chunked prefill：整个 prompt 范围内（排除保护区）Top‑K 需要保留的数量。
- `kv_compression_prompt_topk_keep_max`：当前 batch 内 `prompt_topk_keep` 的最大值（常用于后续 kernel/临时 buffer 的最大 K 尺寸）。
**风险**：中（runner 关键路径增加 CPU 侧元数据计算与额外拷贝；默认关闭不影响现有行为）。

## PR10：feat: kvpress runner 支持 chunked prefill prompt-end 一次性 KV compaction
**目的**：补齐 chunked prefill（scheme 3）路径下的“prompt 结束后一次性压缩”执行：从 forward_context 取出 prompt-end 的 Top‑K 索引/长度，并在第一次 decode 前对所有层 KV cache 做前移 compaction。  
**范围**：`vllm/v1/worker/gpu_model_runner.py`  
**内容**：
- 将 `_stash_kv_compression_prompt_payload` / `_maybe_apply_kv_compression_prompt_compaction` 下沉到 `GPUModelRunnerBase`，让 Base/MTP 两种 runner 路径共用。
- 在 `execute_model` 中补齐调用点：
  - forward 结束后调用 `_stash_kv_compression_prompt_payload()`：把 `forward_context._kv_compression_prompt_payload` 写入 request state（`kv_compression_prompt_idx_sorted/keep_len/prompt_len`）。
  - forward 前调用 `_maybe_apply_kv_compression_prompt_compaction()`：在第一次 decode step 之前，对所有 attention layers 的 KV cache 进行一次性前移压缩（`front_compact_inplace_fa_triton`）。
- compaction 完成后清理 request 上的 pending 状态，避免重复执行。
**实现细节（两个函数分别做什么）**：
- `_stash_kv_compression_prompt_payload`（“保存方案”）：
  - 触发条件：`VLLM_ENABLE_KV_COMPRESSION=1` 且 `chunked_prefill_enabled=True`。
  - 输入来源：读取 `get_forward_context()` 中的 `forward_context._kv_compression_prompt_payload`（由后端在“最后一个 prompt chunk 的 forward”期间写入，forward 结束后会被清空/覆盖，不能跨 step 使用）。
  - payload 字段：`req_indices`（batch 行号）、`idx_sorted`（每个请求的 Top‑K 索引序列/排序结果）、`keep_len`（每个请求最终保留的 prompt token 数）、`prompt_lens`（prompt 总长）。
  - 输出落点：把上述字段写入每个请求的 `CachedRequestState`：
    - `rs.kv_compression_prompt_idx_sorted = idx_sorted[i]`
    - `rs.kv_compression_prompt_keep_len = keep_len[i]`
    - `rs.kv_compression_prompt_prompt_len = prompt_lens[i]`
- `_maybe_apply_kv_compression_prompt_compaction`（“执行方案”）：
  - 触发条件：同样要求 KV compression + chunked prefill；并且仅对已 stash 且已进入 decode 阶段的请求生效（`rs.kv_compression_prompt_idx_sorted != None` 且 `rs.num_computed_tokens >= rs.num_prompt_tokens`）。
  - 执行动作：
    - 收集所有 pending 请求的 `(req_id, idx_sorted, keep_len)`，组成 batch；
    - 构造 `idx_batch[B, K_max]` 与 `keep_tensor[B]`（把每个请求的索引序列 pad 到统一的 `K_max`）；
    - 对每个 KV cache group，拼出该 group 的 `block_table[B, max_blocks]`，并对该 group 内每个 attention layer 的 KV cache 调用 `front_compact_inplace_fa_triton` 做 in-place 前移压缩；
    - compaction 成功后清空 `rs.kv_compression_prompt_idx_sorted/keep_len/prompt_len`，避免重复执行。
**依赖/前置**：需要后续 backend（FlashAttention）在最后一个 prompt chunk 生成并写入 `forward_context._kv_compression_prompt_payload`；仅有本 PR 不会自动产生 payload。  
**兼容性备注**：当 `VLLM_ENABLE_TBO=1` 时，TBO 执行路径目前不会在其内部调用 stash，因此 scheme 3 的 prompt-end payload 可能丢失。本 PR 在 `KV compression + chunked prefill` 场景下会自动绕过 TBO（输出 warning_once），走常规执行路径以保证正确性。  
**风险**：中（runner 关键路径引入一次性 KV 前移；默认关闭不影响现有行为）。
**待处理问题**：
- `_maybe_apply_kv_compression_prompt_compaction` 调用 `self._extract_layer_index(layer_name)`，但 `GPUModelRunnerBase` 并未定义该方法；在 chunked prefill 进入首次 decode 时会触发 `AttributeError`。建议使用已有 `extract_layer_index` helper 或缓存 layer_name→index 映射。
- 共享 KV cache 场景下可能对同一底层 KV 张量重复 compaction（按 layer_name 循环），导致基于原布局的索引被覆盖、KV 被破坏。建议跳过 `shared_kv_cache_layers` 或按 KV cache 实例去重。

## PR11：feat: kvpress flash_attn 透传 KV 压缩元数据
**目的**：把 runner 侧生成的 Top‑K KV 压缩元数据透传到 FlashAttention backend（`FlashAttentionMetadata`），为后续的打分/Top‑K 选择/compaction 执行提供输入。  
**范围**：`vllm/v1/attention/backends/flash_attn.py`  
**内容**：
- 扩展 `FlashAttentionMetadata`：新增 `kv_compression_must_keep/topk_budget/topk_budget_max`，以及 chunked prefill（scheme 3）使用的 `kv_compression_prompt_end/prompt_lens/prompt_topk_keep/prompt_topk_keep_max`。
- 在 `FlashAttentionMetadataBuilder.build()` 中按需从 runner 侧 buffer 挂载这些字段：
  - 非 chunked prefill：仅当 `runner.kv_compression_needs_compaction=True` 时透传 per-token/per-request 元数据（避免无意义开销）。
  - chunked prefill：仅当本 batch 存在 `prompt_end=True` 的请求时透传 prompt-end 元数据（用于后续一次性压缩准备）。
- `topk_budget_max` 通过读取 runner 的 CPU staging buffer 计算，避免 device→host 同步。
**备注**：本 PR 仅“数据打通”，不包含实际的 SnapKV 打分、Top‑K 选择或写回 compaction（这些在后续 PR 完成）。  
**风险**：低（默认关闭不影响；开启时仅增加 metadata 透传）。
**待处理问题**：
- 多 KV cache group 场景下，`_kv_compression_compact_slots` 在 forward context 里按 layer 共享复用会把不同 group 的 `block_table` 混用，导致 compaction 写入错误 block、KV cache 被破坏。建议按 KV cache group 缓存/重算 `dst`，避免跨 group 复用 slot mapping。
- KV cache 跨层共享时，当前 compaction 会在拥有 cache 的层执行后立即写回，可能覆盖共享层在同一步仍将读取的旧布局，导致错误 KV 读取。建议在共享 cache 场景跳过 compaction，或延后到所有共享该 cache 的层执行完后再做。

## PR12a：feat: kvpress flash_attn（scheme 3）生成 prompt-end payload
**目的**：在 chunked prefill（scheme 3）下，仅在“最后一个 prompt chunk”的 forward 期间计算全 prompt 的 Top‑K 保留索引，并写入 `forward_context`，供 runner 在下一步（第一次 decode 前）做一次性 KV 前移 compaction。  
**范围**：`vllm/v1/attention/backends/flash_attn.py`  
**内容**：
- 在 `FlashAttentionImpl.forward()` 中检测 `kv_compression_prompt_end/prompt_lens/prompt_topk_keep` 元数据；若本 batch 存在 `prompt_end=True` 的请求，则调用 `_compute_prompt_end_indices(...)` 计算 payload，并写入 `forward_context._kv_compression_prompt_payload`（同一 step 内只生成一次）。
- `_compute_prompt_end_indices(...)` 的主要流程：
  - 从 `query_start_loc` 切出被选中的请求，并构造 packed 的 query window（每个请求最后 `VLLM_KV_COMPRESSION_SNAPKV_WINDOW` 个 query）。
  - 使用 `gather_k_to_packed_triton` 按 `block_table + prompt_lens` 从 KV cache gather 出完整 prompt keys（packed 形式）。
  - 使用 `query_aware_key_scores`（Triton）计算每个 prompt key 的 token score（token-shared：跨 KV heads sum）；异常时 fallback 到 PyTorch 参考实现。
  - `token_scores` 做 TP all-reduce，确保各 rank 选择一致。
  - 调用 `_prompt_end_topk_keep_indices(...)` 根据 `protected_prefix/suffix/keep_last` 生成 must-keep，并在候选区做 Top‑K，输出 `idx_sorted`（升序）与 `keep_len`。
- payload 字段：`req_indices / idx_sorted / keep_len / prompt_lens`；runner（PR10）会 stash 并在 decode 前对各层 KV 做一次性 compaction。
**风险**：中（最后一个 prompt chunk 会额外做一次打分 + 索引选择；默认关闭不影响现有路径）。
**待优化 / TODO**：
- PyTorch fallback 的打分（reference）在 scheme 3（`_compute_prompt_end_indices`）与非 chunked（`_snapkv_like_token_scores`）会各有一份实现，存在重复维护风险；后续可抽成共享 helper，并补充一致性测试。

## PR12b（拆分）：feat: kvpress flash_attn 实现非 chunked Top‑K compaction（核心）
**目的**：在非 chunked prefill 下，根据 runner 透传的 `must_keep/topk_budget` 做 SnapKV-like 打分与 Top‑K 选择，并在 step 结束后对新写入的 KV 做“前移重写”（drop + pack）。  
**范围**：`vllm/v1/attention/backends/flash_attn.py`  
**内容要点（计划拆分提交）**：
- `token_scores`：仅当 `topk_budget_max>0` 时计算（mixed batch 优化：budget=0 的请求不参与打分）。
- `dst_slots`：用 `_topk_kv_compact_slot_mapping(...)` 生成每个 token 的目标 KV slot（drop 为 -1），保证与 `num_kv_tokens` 的账本一致。
- `writeback`：调用 `reshape_and_cache_*` 把需要保留的 KV 重写到 `dst_slots`；可选缓存 `dst_slots`（跨层共享或 per-layer）。
**风险**：中/高（改动 attention backend 的写回路径；默认关闭不影响）。

## PR12c（拆分）：perf: kvpress flash_attn 支持 async writeback（可选）暂缓/未实现（本轮回退，不进入暂存区）
**目的**：将 compaction 的 `reshape_and_cache_*` 写回放到独立 CUDA stream，和默认 stream 的计算重叠；下一次读写该 layer KV cache 前用 event 同步，保证 `num_kv_tokens` 推进语义正确。  
**范围**：`vllm/v1/attention/backends/flash_attn.py`  
**风险**：中（CUDA stream/event 语义与 cudagraph 兼容性需谨慎；默认关闭不影响）。
**状态**：暂缓/未实现（本轮回退，不进入暂存区）。

## PR13：refactor: 抽取调度器 KV 长度账本逻辑；修复 num_kv_tokens 初始化
**commit**：`3bc7eb74307e954036ddb4fa1fbfe86b81dbca49`  
**目的**：把调度器侧 `num_kv_tokens` 的初始化/推进逻辑收拢到独立模块，减少 scheduler 主流程复杂度；同时补齐 KV compression 开启时 WAITING→RUNNING 但已存在 cached/computed tokens 的 `num_kv_tokens` 初始化，避免后续 KV 写入 offset 错位。  
**范围**：`vllm/v1/core/sched/scheduler.py`、`vllm/v1/kv_compression/scheduler_accounting.py`  
**内容**：
- 新增 `maybe_init_num_kv_tokens_on_running_transition(...)`：RUNNING 过渡时按需初始化 `num_kv_tokens`。
- 新增 `update_num_kv_tokens_after_schedule(...)`：按 chunked prefill / 非 chunked 规则统一推进 `num_kv_tokens`。
- `scheduler.py` 改为调用上述 helper，替代原本内联的大段账本更新逻辑。  
**风险**：低/中（默认关闭不影响；开启 KV compression 时修正边界场景的正确性）。

## PR14：refactor: 将 kv_compression 的 Triton 内核迁移到 vllm/v1/kv_compression
**目的**：把 KV compression 相关的 Triton 内核从 `vllm/v1/attention/kv_compression/` 归位到 `vllm/v1/kv_compression/`，避免模块散落在 attention 目录下，便于后续复用/单测与维护。  
**范围**：`vllm/v1/attention/kv_compression/kv_cache_triton.py`、`vllm/v1/attention/kv_compression/snapkv_triton.py` → `vllm/v1/kv_compression/`  
**内容**：
- 迁移 `kv_cache_triton.py`（KV cache gather / 前移 compaction Triton 内核）。
- 迁移 `snapkv_triton.py`（SnapKV-like 打分 Triton 内核）。
- 删除旧目录下空的 `__init__.py`。  
**待处理问题**：
- 仓库内仍有旧 import：`benchmarks/kvpress/gpu_model_runner.py:3341`、`benchmarks/kvpress/gpu_model_runner.py:3837` 仍从 `vllm.v1.attention.kv_compression.kv_cache_triton` 导入，触发 KV compaction 代码路径会 `ImportError`；需更新为 `vllm.v1.kv_compression.kv_cache_triton`，或在旧路径添加兼容 shim（re-export）。
**风险**：低/中（文件位置调整；需要处理旧路径兼容/调用点更新）。

## PR15：refactor: 抽离 flash_attn 的 KV compression 逻辑到 vllm/v1/kv_compression
**目的**：将 FlashAttention backend 中与 KV compression 相关的打分/Top‑K 选择/slot mapping/forward_context glue 等实现从 `flash_attn.py` 抽离到 `vllm/v1/kv_compression/`，减少后端文件复杂度，便于后续复用（其它 backend）与补单测。  
**范围**：`vllm/v1/attention/backends/flash_attn.py`、`vllm/v1/kv_compression/*`  
**内容**：
- 新增/归位模块：`flash_attn_hooks.py`、`metadata.py`、`snapkv_score.py`、`topk_select.py`、`slot_mapping.py`、`prompt_end.py`、`compaction_step.py`、`forward_context.py`、`kv_cache_view.py`。
- `flash_attn.py` 改为调用 `maybe_compute_prompt_end_payload_flash_attn(...)` 与 `maybe_compact_kv_cache_flash_attn(...)` 等 hooks，避免在 backend 内部维护大段 KV compression glue 代码。
- 统一 chunked prefill（scheme 3）与非 chunked 的关键复用逻辑：packed varlen 坐标、Top‑K keep mask/local rank 计算、dst slot 重写过滤等。  
**风险**：中（重构涉及 attention backend 关键路径；默认关闭不影响现有行为）。
**待处理问题**：
- [P1] 当 `num_kv_heads == block_size` 时需消除 KV cache 布局歧义（`vllm/v1/kv_compression/kv_cache_view.py:26-32`）：`paged_k_cache_view_for_triton_gather` 通过 `key_cache.shape[1] == block_size` 推断布局并 permute 到 HND；但在 ROCm 上 key cache 本身可能已是 `[num_blocks, H, block_size, D]`，若 `H == block_size`（或外部 connector 暴露 HND 且满足该相等关系）会误判并二次 permute，导致 prompt-end scoring 的 Triton gather 以错误的 head/token stride 读取 key，进而产生错误 SnapKV 分数与错误的 prompt-end compaction 索引。建议优先用 `current_platform.is_rocm()` 作为主要判别条件，或显式处理 `shape[1] == shape[2] == block_size` 的歧义情况。