feat: 实现 H2O 逐 head 累计打分与原生 KV 压缩 - #29
Open
kuma-loong wants to merge 11 commits into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
问题与改动
原 H2O 为所有 query heads 保存一份共享 token 分数和保留索引,无法表达 MHA 的逐 head 独立选择。本 PR 改为逐 query head 累计 FP32 attention probability,并在驱逐时确定保留项:
与主线配置对齐
prefill_sparse_method=h2o_prefill控制中间 chunk 压缩至h2o_prefill_budget。sparse_method=h2o控制最终 prompt 压缩至h2o_decode_budget。另补齐当前 FlashInfer paged decode provider 的 head dimension 限制,避免进入不支持该维度的 split-KV merge 路径后才失败。
验证
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: