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

docs/en/features/sparse-methods.md

@@ -49,16 +49,21 @@ compatibility rule.
 SnapKV defaults `sparse_prefill_score_mode` to `logits`; `probability` remains
 an explicit reproducibility option because its additional normalized QK sweep
 is substantially more expensive in measured long-context prefill. PyramidKV
-and H2O continue to default to `probability`. For the shared H2O prompt-scoring
-state this is the canonical path: every KV layer independently sums its
-normalized softmax attention probabilities over the full current query chunk,
-then accumulates that attention mass across prefill chunks. Decode score
-collection and eviction are intentionally disabled. Sparse-vLLM
-reuses FA3's softmax LSE and performs one additional QK sweep because FlashAttention
-does not materialize its probability matrix. `h2o_prefill_score_window=0` selects
-the full current chunk and is the canonical default. A nonzero window in `[1, 128]`
-or explicit `logits` mode is a non-canonical approximation; neither changes the
-requirement that every H2O KV layer computes and retains its own prefill score.
+and H2O use `probability`. H2O accumulates FP32 probability sums per query
+head across prefill chunks. At eviction, max combines the cumulative scores
+within each GQA KV group, or across the whole layer for MLA. MHA heads select independently. Each selection retains
+heavy hitters plus recent tokens within its token budget; native GQA KV sharing
+and MLA latent storage are preserved. MLA H2O prefill currently requires TP1.
+
+`h2o_prefill_score_window=0` scores all queries in each chunk. Windows in
+`[1, 128]` are explicit approximations. H2O rejects `logits` mode because a
+reduced logit vector cannot represent per-head cumulative probabilities.
+When attention provides its softmax LSE, scoring reuses it; otherwise it
+recomputes normalization from the same visible keys. Intermediate chunk eviction
+changes subsequent attention, so results can depend on chunk size and budgets.
+Decode scoring and eviction remain disabled. With `sparse_method=h2o`,
+`h2o_decode_budget` determines final-prefill retention; prefill-only H2O skips
+this final compaction. The cache then grows during generation.
 
 ## Prefill Scheduling Policies
 

docs/zh/features/sparse-methods.md

@@ -38,15 +38,19 @@ prefill attention 计算。它们是同一条轴上的备选项，可以分别
 
 SnapKV 的 `sparse_prefill_score_mode` 默认值改为 `logits`；`probability`
 仍可显式启用以复现实验，但它需要额外执行归一化 QK sweep，在已测长上下文
-prefill 中开销明显更高。PyramidKV 和 H2O 继续默认使用 `probability`。对两个阶段
-共享的 H2O prompt scoring state 而言，这是 canonical 路径：每个 KV layer 都独立地对
-完整当前 query chunk 的归一化 softmax attention probability 求和，并在
-prefill chunk 之间累计 attention mass。当前明确关闭 decode score 收集与
-淘汰。Sparse-vLLM 复用 FA3 的
-softmax LSE；由于 FlashAttention 不物化 probability matrix，还需额外执行
-一遍 QK。`h2o_prefill_score_window=0` 表示完整当前 chunk，是 canonical
-默认设置；`[1, 128]` 的非零 window 或显式 `logits` 模式均属于非 canonical
-近似，但都不会改变每个 H2O KV layer 必须独立计算并保存 prefill score 的要求。
+prefill 中开销明显更高。PyramidKV 和 H2O 使用 `probability`。H2O 逐 query head 跨 prefill chunks
+累计 FP32 概率和；驱逐时通过 max，
+在 GQA 的每个 KV 组内或 MLA 的整个 layer 内归约累计分数。MHA 各 head 独立选择。
+每套选择在预算内保留 heavy hitters 和 recent tokens，保持 GQA 的原生 KV 共享
+以及 MLA 的原生 latent 存储。MLA H2O prefill 当前要求 TP1。
+
+`h2o_prefill_score_window=0` 观察完整当前 chunk；`[1, 128]` 的窗口属于显式近似。
+H2O 拒绝 `logits` 模式，因为归约后的 logits 无法表示逐 head 累计概率。
+Attention 提供 softmax LSE 时复用该结果，否则使用同一可见 KV 集合重新计算归一化。
+中间 chunk 的实际驱逐会改变后续 attention，结果因此可能随 chunk size 和预算变化。
+Decode 评分和驱逐仍保持关闭；仅当 `sparse_method=h2o` 时，
+`h2o_decode_budget` 用于最后一个 prefill chunk 的保留预算。仅启用 H2O prefill 时不做这次最终压缩；
+之后缓存随生成增长。
 
 ## Prefill Scheduling Policy
 

src/sparsevllm/configs/sparse.py

@@ -148,6 +148,8 @@ def _normalize_snapkv(config) -> None:
 
 
 def _normalize_h2o(config) -> None:
+    if getattr(config, "sparse_prefill_score_mode", "probability") != "probability":
+        raise ValueError("H2O per-head accumulation requires sparse_prefill_score_mode='probability'.")
     _normalize_positive_int(config, "h2o_decode_budget", fallback=0)
     _normalize_positive_int(config, "h2o_decode_eviction_interval", fallback=0)
     _normalize_int_attr(config, "h2o_prefill_budget", fallback=0)
@@ -162,14 +164,7 @@ def _normalize_h2o(config) -> None:
             f"h2o_recent_ratio must be in (0, 1), got {config.h2o_recent_ratio}."
         )
     _normalize_int_attr(config, "h2o_prefill_score_window", fallback=0)
-    score_mode = getattr(config, "sparse_prefill_score_mode", "probability")
-    if score_mode == "logits":
-        if config.h2o_prefill_score_window < 0:
-            raise ValueError(
-                "h2o_prefill_score_window must be non-negative in logits "
-                f"mode (0 means the full chunk), got {config.h2o_prefill_score_window}."
-            )
-    elif not 0 <= config.h2o_prefill_score_window <= 128:
+    if not 0 <= config.h2o_prefill_score_window <= 128:
         raise ValueError(
             "h2o_prefill_score_window must be in [0, 128] in probability mode "
             "(0 means the full current chunk), got "