"tests/kernels/test_activation.py" did not exist on "897cb2ae28e93de1b22ecfbffcccfb9493f8f4d9"
Commit ac13c955 authored by laibao's avatar laibao
Browse files

docs: 新增 KVPress v2 PR 整理文档

- 新增 pr/kvpress_v2.md 文档,整理 KVPress v2 相关的 15 个 PR
- 包含各 PR 的目的、范围、内容、风险和待处理问题
- 涵盖 KV 压缩配置、预算计算、Triton 内核、调度层逻辑、runner 侧实现等模块
parent 9d4330d2
# 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` 的歧义情况。
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment