Skip to content

feat: 实现 H2O 逐 head 累计打分与原生 KV 压缩 - #29

Open
kuma-loong wants to merge 11 commits into
CURRENTF:mainfrom
kuma-loong:pr/h2o-per-head-prefill
Open

feat: 实现 H2O 逐 head 累计打分与原生 KV 压缩#29
kuma-loong wants to merge 11 commits into
CURRENTF:mainfrom
kuma-loong:pr/h2o-per-head-prefill

Conversation

@kuma-loong

@kuma-loong kuma-loong commented Sep 7, 2026

Copy link
Copy Markdown
Collaborator

问题与改动

原 H2O 为所有 query heads 保存一份共享 token 分数和保留索引,无法表达 MHA 的逐 head 独立选择。本 PR 改为逐 query head 累计 FP32 attention probability,并在驱逐时确定保留项:

  • MHA 各 head 独立选择;GQA 对同一 KV 组内的累计分数取 max;MLA 对整层取 max,当前要求 TP1。
  • 保留原生 KV head 数与 MLA latent 存储。不同 KV head 的保留项打包到同一组物理槽位,同步压缩各 query head 的分数历史和原始位置;使用临时工作区避免重叠搬运覆盖。
  • SparseController/H2ORuntime 负责打分与逻辑选择,CacheManager 负责物理压缩、槽位释放和状态生命周期。共享 attention 层通过通用接口转发。
  • 移除旧的共享 prefill 选择路径,聚合固定为 max,不提供 mean 配置。容量计算沿用统一 profiling 与 Base/SnapKV 逻辑。

与主线配置对齐

  • prefill_sparse_method=h2o_prefill 控制中间 chunk 压缩至 h2o_prefill_budget
  • sparse_method=h2o 控制最终 prompt 压缩至 h2o_decode_budget
  • 两个阶段可以独立启用;仅启用 prefill 时不做最终 decode-budget 压缩。
  • Decode 打分与周期驱逐仍关闭,生成阶段追加 KV。
  • H2O 使用 probability 打分;窗口为 0 时覆盖整个 chunk,内部按 query 分块限制工作区。

另补齐当前 FlashInfer paged decode provider 的 head dimension 限制,避免进入不支持该维度的 split-KV merge 路径后才失败。

验证

  • H2O、缓存生命周期、阶段配置、算子选择与启动容量相关 CPU 回归:272 passed、19 skipped。
  • CUDA 验证覆盖 MHA/GQA 跨 chunk 数值、逐 head 概率、压缩后原生 decode 读取、MLA 保留及缓存生命周期;修正一处旧测试参数后,所选 37 项均通过。
  • 主线 GLM 模型相关 CPU 测试:28 passed。
  • 新增独立阶段组合的实际缓存保留测试,覆盖 prefill-only、decode-only、同时开启及 FlashPrefill 组合。
  • git diff --check 通过。

本次 rebase 后未重跑完整模型质量评测。逐 head 打分与 KV 搬运的性能优化,以及 profiling 对带历史 chunk 峰值的覆盖,留待后续处理。

MLA TP 支持范围与后续 TODO

本 PR 暂时保留 MLA H2O 仅支持 TP1 的限制;该限制不针对 MHA/GQA。当前逐 head 改造的多卡行为尚未完成验证。

修改前 MLA H2O 没有 TP1 限制,也没有跨 rank 聚合打分:各 rank 根据本地 query heads 的分数独立选择,可能保留不同的原始 token 集合。这种行为并非天然不可行,但需要确认本地 latent KV、位置映射、缓存生命周期和跨 rank 调度始终一致地处理这些差异。因此,本 PR 不直接放开该路径,也不将跨 rank 统一选择视为 MLA TP 的必要条件。

后续 TODO:

  • 明确 MLA TP 的选择语义,优先评估各 rank 内累计、取 max 并独立保留本地 KV;跨 rank 统一选择作为另一种方案单独讨论。
  • 验证各 rank 保留不同 token 集合时,原始位置、物理槽位、请求长度及调度/释放行为是否正确。
  • 补充 TP2 的 chunk prefill、最终压缩、原生 decode 读取、CUDA Graph 和缓存生命周期验证,以匹配所选语义的独立参考结果为准。
  • 完成多卡验证后再移除 MLA TP1 限制,并更新支持范围文档。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant