feat: add topology-aware sparse MLA candidate policy

[teerthsharma]

Summary

I made this a code/test/build-only stacked follow-up to #1160 for topology-aware sparse MLA candidate rewriting.

This PR is intentionally draft until #1160 lands. It includes no docs, README, plan, or spec files in the remote diff.

What changed

• Added topology_kv_policy.py for opt-in topology-aware candidate merge modes.
• Kept default behavior disabled with RTP_LLM_TOPOLOGY_KV_POLICY=disabled.
• Applied the policy only to ragged and CP prefill top-k results.
• Left decode _get_topk_paged unchanged because that path returns page-table / physical block IDs, not token offsets.
• Added absolute-coordinate validation with production fallback on mismatch.
• Normalized policy and coordinate-mismatch env/config strings case-insensitively.
• Restricted policy input indices to torch.int32 / torch.int64, so -1 padding and absolute offsets remain representable.
• Added separate Python-path guards for total policy tokens, topk_indices.numel(), and top-k width.
• Made topology_only consume the full structural budget so available structural tokens fill min(topk_width, length) instead of being capped by merge-mode structural fraction.
• Moved optional drift scores onto the same detached CPU policy path before structural selection.
• Added CUDA-sync and long-prefill guards so the Python policy path cannot silently run on hot CUDA paths.
• Added learned-retention counters and capped structural injection by budget fraction for merge modes.
• Removed shared per-layer fingerprint/counter state; compression-hit semantics now require caller-provided per-request fingerprint context.
• Added Bazel runfiles declarations for benchmark:topology_kv_policy_test and its CUDA manual variant.

Policy modes

• disabled: passthrough, default.
• topology_sparse_merge: merge topology sink/local/witness candidates with learned top-k candidates.
• topology_compress_sparse: sparse merge plus stable scaffold/output-contract counters.
• topology_only: topology candidates only, using the full structural budget.

Coordinate contract

apply_topology_kv_policy requires learned non-padding top-k indices to already be absolute token indices in the downstream sparse-attention coordinate system.

For ragged and CP prefill rows:

valid row interval = [topk_indices_offset, topk_indices_offset + length)

row_starts is only a row-count alignment guard here. It is not added to returned top-k coordinates.

Decode is intentionally bypassed because _get_topk_paged returns page-table / physical block IDs, not token-space offsets.

Validation I ran

Windows / local Python:

• python -m unittest benchmark.test_topology_kv_policy -v: 35 tests run, 34 passed, 1 CUDA-manual test skipped by default.
• python -m unittest benchmark.test_topology_kv_candidate_schedule -v: 22 tests run, 21 passed, 1 CUDA-manual test skipped by default.
• Static BUILD/runfiles check for benchmark:topology_kv_policy_test: passed.
• python -m py_compile ...: passed.
• python -m black --check ...: passed.

Docker / Linux GPU image:

• python -m unittest benchmark.test_topology_kv_policy -v: 35 tests run, 34 passed, 1 CUDA-manual test skipped by default.

I cannot run Bazel in this Windows shell because neither bazel nor bazelisk is installed here. I also cannot run final H20 validation because I do not own that GPU.

Dependency

Blocked by #1160. This branch remains stacked and should be rebased after #1160 lands.

[LLLLKKKK]

AI Code Review - PR #1162

Status: BLOCKING

Summary: P0/0 · P1/1 · P2/3 · P3/1

Blocking Issues

P1

• 生产接线缺集成测试且 topk 坐标系假设未验证，启用后可能静默产生错误索引 @ rtp_llm/models_py/modules/hybrid/indexer.py:173
  • 建议：补一个经由 Indexer._compute_topk（或 mock indexer_op._get_topk_* 返回已知坐标约定的张量）驱动 decode / ragged / cp 三条路径的集成测试，断言合并后索引落在合法范围且与下游 sparse attention 消费约定一致；并在 apply_topology_kv_policy 文档/断言中显式声明 learned topk 的坐标系契约（绝对 vs 行内相对、是否含 offset），对越界值 fail-fast 而非静默丢弃。

Non-blocking Suggestions

P2

• 策略在 decode 热路径上使用 Python 逐行循环与多次 GPU→CPU 同步 @ rtp_llm/models_py/modules/hybrid/topology_kv_policy.py:130
  • 建议：将结构块选择与合并改写为张量化实现（在目标 device 上用 torch 向量操作/散射完成 sink/local/witness 与 learned 的去重合并），避免 .item() / .cpu().tolist() 的逐 step 同步与 Python 循环；计数器统计也尽量批量化或在关闭统计时跳过。至少应给出启用态的 decode 吞吐/延迟基准数据，证明其可用。
• policy 字符串分发未使用 Enum/Literal @ rtp_llm/models_py/modules/hybrid/topology_kv_policy.py:9
  • 建议：用 enum.Enum 或 typing.Literal["disabled","topology_sparse_merge","topology_compress_sparse","topology_only"] 定义策略集合，构造 TopologyKvPolicyConfig 时校验并对未知值 fail-fast（或明确 warn），环境变量解析处集中做一次映射，避免字符串在多文件重复。
• 计数器与 fingerprint 语义含糊、缺少说明，观测性易误导 @ rtp_llm/models_py/modules/hybrid/topology_kv_policy.py:34
  • 建议：为该模块补 module/函数级 docstring，解释 stable_scaffold/output_contract/fingerprint 的用途与来源；将计数器命名与实现对齐真实语义（如 unselected_tokens 而非 raw_kv_tokens_avoided），或删除当前无实际压缩含义的字段，避免观测指标误导。

P3

• 夹带与本 PR 无关的括号缩进改动且不符合 black @ rtp_llm/models_py/modules/hybrid/indexer.py:180
  • 建议：回退该行缩进改动，保持逻辑变更与格式变更分离；提交前跑 black/pre-commit 以确保格式合规。

Checklist Violations (2 fail / 56 total)

General Principles Checklist

• [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue 生产接线缺集成测试且 topk 坐标系假设未验证，启用后可能静默产生错误索引
_compute_topk 三条路径分别以不同的坐标参数调用策略：decode 路径 _apply_topology_kv_policy(topk_result, fmha_params.expanded_seq_lens) 不传 row_starts/offset；ragged 路径传 fmha_params.ks / fmha_params.topk_indices_offset；cp 路径传 cp_params.precomputed_ks / precomputed_topk_off。apply_topology_kv_policy 内部对 structural token 做 token + row_start + offset，并按 [row_start+offset, +length) 过滤 learned 值。但 _get_topk_paged/_get_topk_ragged 返回的 learned 索引究竟是绝对位置还是行内相对、是否已含 offset，未有测试验证。若坐标系不一致，越界值被过滤成 -1 或指向错误 K
• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue 生产接线缺集成测试且 topk 坐标系假设未验证，启用后可能静默产生错误索引
_compute_topk 三条路径分别以不同的坐标参数调用策略：decode 路径 _apply_topology_kv_policy(topk_result, fmha_params.expanded_seq_lens) 不传 row_starts/offset；ragged 路径传 fmha_params.ks / fmha_params.topk_indices_offset；cp 路径传 cp_params.precomputed_ks / precomputed_topk_off。apply_topology_kv_policy 内部对 structural token 做 token + row_start + offset，并按 [row_start+offset, +length) 过滤 learned 值。但 _get_topk_paged/_get_topk_ragged 返回的 learned 索引究竟是绝对位置还是行内相对、是否已含 offset，未有测试验证。若坐标系不一致，越界值被过滤成 -1 或指向错误 K

Strengths

• 特性默认 disabled，通过环境变量门控，对现网默认部署零影响，回滚路径清晰。
• topology_kv_candidate_schedule.py / topology_kv_policy.py 的纯函数单元有较全的单测覆盖（padding、去重、边界、2D/3D 输入、fp16 用 fp32 累加、RNG 状态不被污染等）。
• build_key_block_centroids 对 half precision 采用 fp32 累加，避免精度损失，并有对应测试。
• 输入校验较完整（block_size/seq_len 非负、候选索引形状/重复/越界检查）。

[Copilot]

Pull request overview

Note

Copilot couldn't run its full agentic review because no GitHub Actions runner was available. Make sure your repository has a runner available to run Copilot's review, or add a copilot-setup-steps.yml file specifying one with the runs-on attribute. See the docs for more details.

Adds a topology-aware KV candidate rewrite policy (with optional compression counters/fingerprinting) and wires it into the Hybrid Indexer behind an env-gated switch to preserve existing learned sparse top-k behavior by default.

Changes:

• Introduces topology_kv_policy.py implementing topology-based candidate selection/merge modes plus stable fingerprint + counters.
• Integrates policy application into Indexer._compute_topk behind RTP_LLM_TOPOLOGY_KV_POLICY=disabled default.
• Adds manual-only benchmark helpers and unit tests for topology scheduling and the policy helper.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
rtp_llm/models_py/modules/hybrid/topology_kv_policy.py	New helper implementing topology candidate merge/only modes plus compression counters/fingerprint.
rtp_llm/models_py/modules/hybrid/indexer.py	Env-gated integration hook to rewrite top-k indices and expose latest counters/fingerprint.
benchmark/topology_kv_candidate_schedule.py	Benchmark utility to build topology candidate schedules and compare dense vs sparse decode attention.
benchmark/test_topology_kv_policy.py	Tests for policy behavior, counters/fingerprint, and a CUDA sparse-attention E2E check.
benchmark/test_topology_kv_candidate_schedule.py	Tests for candidate schedule construction, validation, and benchmark determinism.
benchmark/BUILD	Adds manual Bazel py_test targets to avoid entering wildcard GPU CI.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[LLLLKKKK]

AI Code Review - PR #1162

Status: BLOCKING

Summary: P0/0 · P1/0 · P2/4 · P3/0

Non-blocking Suggestions

P2

• 全局 _TOPOLOGY_KV_STATE 跨请求/并发共享导致 fingerprint 与 counters 污染 @ rtp_llm/models_py/modules/hybrid/indexer.py:20
  • 建议：明确该状态的作用域：若仅为观测用途，建议在文档/注释中标注「仅反映最近一次调用、不保证请求隔离、非并发安全」；若希望 compression_hits 有意义，应改为按请求/序列维度携带 fingerprint（例如通过调用方传入 per-request 上下文），而非模块级全局共享。
• 开启 merge 后 top-k 顺序重排并驱逐部分 learned token，真实稀疏内核路径缺自动数值覆盖 @ rtp_llm/models_py/modules/hybrid/topology_kv_policy.py:2217
  • 建议：补充一个不依赖 GPU 的确定性测试，断言 merge 输出满足生产内核对 index 布局的契约（如 -1 padding 只在行尾、有效 index 唯一且在合法区间）；并在 PR 描述或注释中显式说明该策略会改变注意力语义（驱逐部分 learned token）。若生产 sparse 内核对 index 顺序有要求（例如需升序），应在返回前排序或明确记录内核为顺序无关。
• 在 CUDA 张量上开启策略默认静默 no-op，允许 sync 时于 prefill 热路径引入 device→host 同步与 Python 逐行循环 @ rtp_llm/models_py/modules/hybrid/indexer.py:238
  • 建议：在 CUDA 且未开启 sync 时至少打一条一次性 warning（或启动期 log），说明策略被跳过及原因，避免「开了没效果」的运维困惑（G.6.1.11）；并在文档中标注开启 sync 会带来 host 同步 + Python 逐 token 开销，建议仅用于实验/评测而非在线主路径。
• 测试文件使用 unittest.mock 但未显式导入，存在按导入顺序失败的风险 @ benchmark/test_topology_kv_policy.py:931
  • 建议：在文件顶部显式增加 from unittest import mock（并将两处改为 mock.patch.dict），或直接 import unittest.mock，消除对传递性导入的隐式依赖，保证测试在隔离环境下稳定通过。

Checklist Violations (5 fail / 56 total)

General Principles Checklist

• [6.1] Architecture — 可观测性：日志/指标/超时可操作、非噪声 → issue 在 CUDA 张量上开启策略默认静默 no-op，允许 sync 时于 prefill 热路径引入 device→host 同步与 Python 逐行循环
_apply_topology_kv_policy 中：if topk_result.is_cuda and os.getenv("RTP_LLM_TOPOLOGY_KV_ALLOW_CUDA_SYNC") != "1": return topk_result。运维如果只设置 RTP_LLM_TOPOLOGY_KV_POLICY=topology_sparse_merge（GPU 部署常态）而未额外设置 sync flag，策略会静默不生效、且不更新 counters、无任何日志，难以从行为上察觉「配置未起作用」。反之设置 sync flag 后，apply_topology_kv_policy 会 topk.detach().cpu() 并对最多 max_policy_tokens 行、每行 topk=2048 宽做纯 Python 循环，在 prefill 热路径引入同步与显著 CPU 开销。
• [6.1] Architecture — 状态不变量：创建/更新/失败/重试/回滚路径有效 → issue 全局 _TOPOLOGY_KV_STATE 跨请求/并发共享导致 fingerprint 与 counters 污染
_TOPOLOGY_KV_STATE = weakref.WeakKeyDictionary() 以 Indexer 模块实例（每层一个，被所有并发请求/batch 共享）为 key。_apply_topology_kv_policy 先读 self.latest_topology_kv_fingerprint（上一次 forward 写入的值），再写 _TOPOLOGY_KV_STATE[self]。由于该状态是「按层」而非「按请求/序列」，previous_fingerprint == fingerprint 判定出的 compression_hits、以及 latest_topology_kv_counters 反映的都只是「最近一次经过该层的任意请求」，对不同请求而言语义无意义；并发 prefill 下 read-modify-write 还存在竞态。
• [6.1] Architecture — 错误语义：fail-fast/retry/fallback/silent 行为显式 → issue 在 CUDA 张量上开启策略默认静默 no-op，允许 sync 时于 prefill 热路径引入 device→host 同步与 Python 逐行循环
_apply_topology_kv_policy 中：if topk_result.is_cuda and os.getenv("RTP_LLM_TOPOLOGY_KV_ALLOW_CUDA_SYNC") != "1": return topk_result。运维如果只设置 RTP_LLM_TOPOLOGY_KV_POLICY=topology_sparse_merge（GPU 部署常态）而未额外设置 sync flag，策略会静默不生效、且不更新 counters、无任何日志，难以从行为上察觉「配置未起作用」。反之设置 sync flag 后，apply_topology_kv_policy 会 topk.detach().cpu() 并对最多 max_policy_tokens 行、每行 topk=2048 宽做纯 Python 循环，在 prefill 热路径引入同步与显著 CPU 开销。
• [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue 开启 merge 后 top-k 顺序重排并驱逐部分 learned token，真实稀疏内核路径缺自动数值覆盖
_merge_row 先放入 structural（local 范围）token 至 structural_budget（默认 budget*0.5），再放 learned token，最终 merged = torch.tensor(rows, ...)。这既改变了输出相对底层 op 的 token 顺序（structural 前置），也会在预算满时驱逐部分 learned top-k（learned_evicted_tokens），从而改变注意力实际关注的 KV 集合。仓库内唯一的真实数值 e2e（test_cuda_policy_output_runs_sparse_decode_attention_e2e）被 skipUnless(manual + cuda) 关闭，且比对使用 benchmark 内基于 index_select 的顺序无关实现，而非生产 sparse MLA 内核；因此「开启策略后生产内核输出仍正确/内核对顺序与 -1 padding 布局的容忍性」没有 CI 可自动验证的覆盖。
• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue 测试文件使用 unittest.mock 但未显式导入，存在按导入顺序失败的风险
test_topology_kv_policy.py 顶部仅 import unittest（未 import unittest.mock 或 from unittest import mock），但第 931、941 行调用 unittest.mock.patch.dict(...)。CPython 的 import unittest 并不会自动导入 unittest.mock 子模块；只有当同进程中其它模块已 import unittest.mock 时该属性才可用。该 target 由 BUILD 中 main=test_topology_kv_policy.py 经 unittest.main() 运行，若运行环境下无人先行导入 unittest.mock，这两个用例（test_indexer_topology_env_int_reports_bad_value / _float）将抛 AttributeError: module 'unittest' has no attribute 'mock'。

Strengths

• 特性默认 disabled，且参数评估的 fmha_params.ks/expanded_seq_lens/topk_indices_offset 已被底层 _get_topk_ragged 使用，禁用态下不引入新的 AttributeError 风险，生产 blast radius 很小。
• 坐标系处理严谨：topk_indices 被明确定义为绝对坐标 [offset, offset+length)，越界时按 coordinate_mismatch_action 选择 raise 或 fallback_disabled，避免静默丢弃掩盖坐标系错配。
• 单测覆盖充分（padding、半精度 fp32 累加、去重、因果性、结构预算、offset 坐标、CP/ragged/decode 三路分派、env 解析错误等），且 CUDA 端 e2e 用 skipUnless + env gate 隔离。
• 异常处理规范（raise X from Y）、字符串分发用 Literal/frozenset 而非裸字符串、dataclass(frozen=True) 承载配置与计数器。

[LLLLKKKK]

AI Code Review - PR #1162

Status: BLOCKING

Summary: P0/0 · P1/1 · P2/1 · P3/0

Blocking Issues

P1

• benchmark 新增的 topology_kv_policy_test target 在 bazel 下导入即失败 @ benchmark/BUILD:19
  • 建议：将 topology_kv_candidate_schedule.py 加入 topology_kv_policy_test（及其 _cuda_manual 变体）的 srcs，并把被测的 //rtp_llm/models_py/modules/hybrid:topology_kv_policy.py 与 indexer.py 通过 data（或改为依赖对应 py_library）声明进 target，使 MODULE_PATH/INDEXER_PATH 在 runfiles 中可解析；或将这两个用例改为依赖已打包的 rtp_llm 库导入而非按文件系统路径读取源码。修复后建议本地/CI 实跑确认 target 通过。

Non-blocking Suggestions

P2

• prefill 应用策略而 decode 直接旁路，启用时 prefill/decode 候选集不一致 @ rtp_llm/models_py/modules/hybrid/indexer.py:272
  • 建议：在 docstring / 环境变量说明中显式声明「策略仅作用于 prefill、decode 保持原始 top-k」这一契约与其动机；或评估是否需要在 decode _get_topk_paged 路径上给出对齐处理（哪怕仅是文档化的限制），避免使用者误以为整条推理链均受策略约束。

Checklist Violations (1 fail / 56 total)

General Principles Checklist

• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue benchmark 新增的 topology_kv_policy_test target 在 bazel 下导入即失败
topology_kv_policy_test 的 srcs 只有 test_topology_kv_policy.py。但该测试文件顶层先 import benchmark.topology_kv_candidate_schedule，回退分支 import topology_kv_candidate_schedule——该文件未被列入本 target 的 srcs，bazel sandbox runfiles 中不存在，会抛未被捕获的 ModuleNotFoundError。此外模块加载还依赖 Path(__file__).resolve().parents[1] / "rtp_llm" / ... / "topology_kv_policy.py" 及同目录 indexer.py 直接读源文件（spec.loader.exec_module、INDEXER_PATH.read_text()），这些源文件同样未声明为 srcs/data，sandbox 内不可见。因此该 target 一旦被执行即在模块收集阶段失败（对比 `topol

Strengths

• 特性默认关闭（RTP_LLM_TOPOLOGY_KV_POLICY=disabled），CUDA 路径二次门控，坐标不匹配默认回退原始 top-k，线上默认路径零影响。
• env 解析（_topology_env_int/_float）与 TopologyKvPolicyConfig.__post_init__ 校验均 fail-fast，错误信息明确。
• topology_kv_policy.py 顶部 docstring 清楚界定了绝对坐标系、topk_indices_offset 的有效区间、以及 counters 的可观测性语义。
• 针对策略与调度逻辑（去重、因果、padding 布局契约、fingerprint 命中）提供了较完整的 CPU 单测覆盖。

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[LLLLKKKK]

AI Code Review - PR #1162

Status: BLOCKING

Summary: P0/0 · P1/1 · P2/1 · P3/0

Blocking Issues

P1

• topology_only 默认只保留一半结构候选 @ rtp_llm/models_py/modules/hybrid/topology_kv_policy.py:422
  • 建议：对 topology_only 使用完整 budget，或调用 _merge_row 时强制 max_structural_fraction=1.0；补一条 topology_only 测试，断言可用结构 token 会填满 min(topk_width, length)。

Non-blocking Suggestions

P2

• policy 接受的 dtype 无法保证 padding 和绝对坐标语义 @ rtp_llm/models_py/modules/hybrid/topology_kv_policy.py:49
  • 建议：将 policy 输入限制为生产实际使用的 torch.int32/torch.int64，或在输出前统一提升到可表示 -1 和最大 offset 的有符号 dtype，并补充 unsigned/narrow dtype 的拒绝测试。

Checklist Violations (2 fail / 47 total)

General Principles Checklist

• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue ``topology_only 默认只保留一半结构候选
  _`topology_only` 把 `learned_values` 置为全 `-1`，但仍走 `merge_row` 的 `structural_budget=int(budget * max_structural_fraction)`。默认 0.5 时，topk=2048 只会输出最多 1024 个结构 token，剩余补 `-1`，即使结构候选充足。
• [6.1] Tests — 边界 case 覆盖（空、单元素、最大值） → issue policy 接受的 dtype 无法保证 padding 和绝对坐标语义
  模块约定负数是 padding，且输出会带 topk_indices_offset 的绝对 token 坐标；但 INTEGER_TOPK_DTYPES 允许 torch.uint8/int8/int16，随后用原 dtype 构造 merged。这些 dtype 不能可靠表达 -1 padding 或较大的 token offset。

Strengths

• 默认 RTP_LLM_TOPOLOGY_KV_POLICY=disabled，CUDA tensor 还需要额外显式允许 host-sync，默认线上路径风险被控制住。
• 新增单测覆盖了 ragged/CP topk_indices_offset 合约、coordinate mismatch fallback、size bypass 和 CUDA manual e2e 入口。