feat: duplicated kv cache.

[ZhangZhiPku]

Add duplicate kv cache logic for num_kv_heads can not be divided by tp_size.

[LLLLKKKK]

Code Review

KV head 不整除 TP 时使用 GCD 复制策略，方向正确。

问题 1: 移除了 std::max(1, ...) 防护

MHAKVCacheSpec.h 旧代码有 std::max(1, ...) 保证每个 rank 至少 1 个 KV head。新代码移除了这个防护，如果 kv_head_num 为 0（配置错误），local_head_num_kv 会变成 0，导致下游除零或零大小分配。

问题 2: 相同的 8 行 KV head 分区注释在 5 个文件中重复

MHAKVCacheSpec.h、ModelConfig.cc、model_weight.py（3 处）都有完全相同的逻辑注释。建议提取为共享工具函数，避免后续修改遗漏。

问题 3: 缺少 duplicated weights 正确性测试

diff 中只改了分区数学，没有测试验证 duplicated KV weights 在多 rank 下产生正确的 attention 输出。

[LLLLKKKK]

🤖 AI Code Review

概述

对 GenerateStream 状态机进行重构：消除 STOPPED/PAUSED 状态，引入 LOADING_CACHE 状态支持异步 KV Cache 加载，分离错误报告（reportError）与状态转换。新增约 1000 行测试代码。

优点

• 状态机简化设计清晰，LOADING_CACHE 状态和 loading_cache_streams_ 队列设计合理
• 错误报告与状态转换分离是好的架构改进
• 测试覆盖充分，覆盖状态机转换、异步加载调度、错误处理等场景

建议改进

P0 - Bug

1. reportError() 在已持有 output_mutex_ 的上下文中调用会导致死锁：reportError() 内部获取 output_mutex_，但以下调用点在已持有该锁的情况下调用它：

   • GenerateStream::update() 中 token 校验失败和 kv cache blocks 更新失败时
   • GenerateStream::specUpdate() 中
   • NormalGenerateStream::enqueueGenerateOutput() 中（通过 update() -> updateOutput() 调用链）

   std::mutex 不支持递归加锁，这将导致死锁。需要提供一个不加锁的 reportErrorWithoutLock() 变体，或将 output_mutex_ 改为 std::recursive_mutex。

P1 - 重要

1. loadCacheDone() 失败时无限重试：connector 持续失败时流将永远卡在 LOADING_CACHE 状态，占用 KV Cache blocks 不释放。建议增加最大重试次数限制。

2. reportError() 替代 setStop() 后部分调用点的流不会被及时终止：PD 分离场景中 DecodeGenerateContextNew 析构时，reportError() 仅标记错误不改变状态，decode engine 可能继续执行无效计算直到下一轮调度。建议确认所有执行路径在关键检查点检查 hasError()。

P2 - 建议

1. loaded_ 标志设置时机过早：在 asyncLoadCache() 返回前就设为 true，建议移到返回之后或重命名为 load_attempted_。
2. evaluateLoadingCacheStreams() 中 setFinishedWithoutLock() 缺少锁保护：scheduler 未持有 stream 的 output_mutex_，而 reportError() 可从任意线程调用，存在竞态。

总结

P0 死锁问题必须修复。建议同时增加 loadCacheDone() 重试上限后再合入。

[LLLLKKKK]

🤖 AI Code Review (修正版 — 之前的评论误贴了其他 PR 的内容，非常抱歉)

概述

当 kv_head_num 无法被 tp_size 整除时（且 kv_head_num != 1），引入 GCD 策略将 KV heads 按 gcd(kv_head_num, tp_size) 分组，使多个 TP rank 共享相同的 KV heads。变更涉及 C++ 配置层、KV cache 规格层和 Python 权重加载层。

优点

• GCD 策略是业界常见做法（vLLM 等也有类似实现）
• 代码改动范围合理，覆盖了配置、cache spec 和权重加载三层

建议改进

P1 - 重要

1. local_head_num % local_kv_head_num 整除关系未校验: GCD 分组后必须满足 head_num % kv_head_num == 0，否则 GQA 分组计算会产生错误结果。建议在 getAttentionConfigs 中增加断言。
2. MHAKVCacheSpec 与 getAttentionConfigs 双路径计算不统一: 两处独立计算 local_kv_head_num，代码中已有 TODO 注释承认此问题。建议抽取为共享静态方法。
3. 移除旧校验后缺少替代性校验: loader.py 删除了 head_num_kv % tp_size != 0 校验但未添加新校验，不合理配置会在后续计算中产生难以调试的错误。
4. 缺少测试覆盖: 核心 TP 权重分配变更无任何单元测试。

P2 - 建议

• std::gcd 需要 #include <numeric>，确认 MHAKVCacheSpec.h 和 ModelConfig.cc 已包含
• LinearKVCacheSpec 未同步更新，逻辑不一致
• Python 侧相同的 KV 分组逻辑重复 5 次，建议提取辅助函数

总结

核心思路正确，但缺少整除关系校验和测试覆盖。建议补充后合入。

[LLLLKKKK]

🤖 AI Code Review — PR #804 — feat: duplicated kv cache.

概述

本 PR 通过 GCD（最大公约数）策略解决了 kv_head_num 不能被 tp_size 整除时的 KV head 分配问题，允许多个 TP ranks 共享（duplicate）相同的 KV heads。改动涉及 C++ 侧的 ModelConfig::getAttentionConfigs 和 MHAKVCacheSpec，以及 Python 侧的三个权重切分函数。

优点

• GCD 策略的数学正确性经验证没有问题：对于所有合理的 (head_num, kv_head_num, tp_size) 组合，local_head_num % local_kv_head_num == 0 始终成立
• 新增了 duplicated_kv_test.py 覆盖多种组合的 Q-to-KV mapping 正确性
• 统一了 kv_head_num == 1 的特殊分支处理（GCD 策略自然覆盖）
• 在 getAttentionConfigs 中新增了 head_num % kv_head_num != 0 的前置校验

建议改进

P1 - 重要

1. GCD 逻辑重复，存在不一致风险
   MHAKVCacheSpec.h 和 ModelConfig.cc 各自独立实现了相同的 GCD 分支逻辑。代码中已有 TODO 注释承认了这个问题。建议提取为一个共享的静态函数（如 computeLocalKvHeadNum(kv_head_num, tp_size)），避免未来修改一处忘记另一处。

2. MHAKVCacheSpec 缺少 head_num % kv_head_num == 0 校验
   ModelConfig::getAttentionConfigs 新增了此校验，但 MHAKVCacheSpec 构造函数没有同步添加。如果有人直接构造 MHAKVCacheSpec 而不经过 getAttentionConfigs，可能传入不合法的配置。

3. splitKVPartitionBytes 的 partition 兼容性
   MHAKVCacheSpec::splitKVPartitionBytes 中有 heads % partition_count == 0 的 assert。在 GCD 策略下 local_head_num_kv 的值可能与之前不同，需要确认所有调用路径中 partition_count 的取值是否兼容。

P2 - 建议

4. LinearKVCacheSpec 未同步更新
   LinearKVCacheSpec.h 仍使用旧的 std::max(1, ...) 模式。如果 qwen3.5 的 linear attention 层也可能出现 linear_num_value_heads % tp_size != 0，这里也需要 GCD 策略。如果保证整除，建议添加注释说明。

5. KV cache 显存影响未说明
   GCD 策略下 local_kv_head_num 可能大于简单除法的结果（如 kv_head=2, tp=8 时每 rank 1 个 head 而非 0.25 个），KV cache 显存会相应增加。建议在 PR description 中补充显存影响分析。

P3 - Nit

6. 测试中 random.randrange 未设置 seed，建议设置固定 seed 以保证可复现性。
7. 三处 KV head partitioning 的 docstring 完全相同，建议提取为模块级常量或引用同一处文档。

总结

GCD 策略数学上正确，权重切分和 KV cache 规格计算保持一致，测试覆盖了核心场景。主要风险是 GCD 逻辑在两处独立实现以及 splitKVPartitionBytes 的 partition 兼容性。建议合并前解决 P1 问题。

[LLLLKKKK]

🤖 AI Code Review — PR #804 — feat: duplicated kv cache. (增量 Review v2)

概述

增量 review，基于新增 commits 8bc90b7 / e7fac91（对比首次 review 基线 e59018e）。新增了 head_num % kv_head_num 前置校验和 160 行单元测试覆盖 Q-to-KV head mapping 正确性。

首次 Review 问题跟踪

#	问题	级别	状态
1	getAttentionConfigs 与 MHAKVCacheSpec GCD 逻辑重复	P1	未修复
2	MHAKVCacheSpec 缺少 head_num % kv_head_num 校验	P1	部分修复（已加在 getAttentionConfigs，MHAKVCacheSpec 仍无 assert）
3	LinearKVCacheSpec 未同步 GCD 逻辑	P2	未修复
4	KV cache 显存影响未说明	P2	未修复
5	splitKVPartitionBytes partition 兼容性	P1	未修复
6	测试 random.randrange 未设 seed	P3	未修复

优点

• 新增测试设计合理，覆盖 4 head_dim x 4 head_config x 4 tp 的组合矩阵，使用 subTest 便于定位
• head_num % kv_head_num 前置校验放在 getAttentionConfigs 是正确位置

建议改进

P1 - 重要（遗留）

1. GCD 逻辑重复：MHAKVCacheSpec.h 和 ModelConfig.cc 仍各自独立实现 GCD 分支。建议提取共享函数。
2. splitKVPartitionBytes 兼容性：GCD 策略下 local_head_num_kv 可能不被 partition_count 整除，需确认所有调用路径。

P2 - 建议

3. 测试覆盖缺口：get_sp_tensor_blocked 未被测试；缺少 kv_head > tp 且不整除的 case（如 kv_head=6, tp=4）；缺少 kv_head=1 退化 case 验证。
4. 测试 BUILD 依赖较重：duplicated_kv_test 依赖了 th_transformer/rtp_compute_ops/fastapi，但测试只用了 torch 和 model_weight，可考虑精简。

P3 - Nit

5. random.randrange 未设 seed，影响可复现性。

总结

增量提交回应了测试覆盖的关切，测试质量不错。但首次 review 的 3 个 P1 中仅 #2 部分修复，#1（逻辑重复）和 #5（partition 兼容性）仍需关注。建议合并前至少确认 splitKVPartitionBytes 在新场景下不会触发 assert。

[LLLLKKKK]

🤖 AI Code Review — PR #804
Head SHA: e7fac91ce4c0e6f28e93461ef8885e667a631ccb | Verdict: P2

Summary

Adds GCD-based KV head partitioning for tensor parallelism when kv_head_num % tp_size != 0. Previously required exact divisibility or kv_head_num == 1. Now uses gcd(kv_head_num, tp_size) to determine per-rank KV head count, allowing multiple TP ranks to share (duplicate) KV heads. Changes span ModelConfig.cc, MHAKVCacheSpec.h, model_weight.py, loader.py, plus a new comprehensive test.

Findings

1. [P2] Duplicated comment blocks (5 copies)
   The same multi-line GCD partitioning explanation is copy-pasted in ModelConfig.cc, MHAKVCacheSpec.h, and 3x in model_weight.py. Extract to a shared doc reference to avoid divergence.

2. [P2] Missing head_num % tp_size validation
   Q heads are split by head_num / tp_size without checking divisibility. Since this PR relaxes TP constraints, an explicit assertion would be prudent.

3. [P2] MHAKVCacheSpec and ModelConfig could diverge
   The TODO comment acknowledges the two independent GCD implementations must stay in sync. Consider having one call the other.

4. [P2] Non-deterministic test
   duplicated_kv_test.py uses random.randrange without a seed. Setting a seed would make failures reproducible.

The GCD approach is mathematically sound and the test is thorough. Good generalization of the previous special cases.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ ZhangZhiPku
❌ kudomcho
_{You have signed the CLA already but the status is still pending? Let us recheck it.}