mimo v2.5 support - #1351
Conversation
LLLLKKKK
left a comment
There was a problem hiding this comment.
AI Code Review - PR #1351
Status: BLOCKING
Summary: P0/0 · P1/4 · P2/17 · P3/8
Reviewed: commit 4d2f754b2099 · 2026-08-29 13:27 UTC+8
Blocking Issues
P1
- SWA 多块前缀复用(swaMatchBegin + CP 槽位映射)无任何单测覆盖 @
rtp_llm/cpp/cache/HybridKVCacheAllocator.cc:65- 建议:在
HybridTypeKVCacheAllocatorTest.cc(已有 SWA policy 脚手架)补reuseCache用例:(1)window_tokens=0与旧单 tail 行为逐位等价;(2) 窗口跨 2~3 个 block 时[begin,pos]全部落位、更早槽位保持 NULL;(3) 窗口内任一 key 未命中时pos逐级回退;(4) 边界值window_tokens=1、window_tokens <= seq_size_per_block、begin被max(0,...)截断至 0。在HybridKVCacheAllocatorCPShardTest.cc补cp_size>1下 compact 与非 compact 两种logical_pos映射断言。建议同时把swaMatchBegin暴露为可直接单测的纯函数入口。
- 建议:在
- 非对称 K/V、attention sink 与窗口换算的数值正确性无 CI 覆盖,唯一新增用例不可证伪 @
rtp_llm/models_py/modules/factory/attention/cuda_impl/test/test_flashinfer_prefill/test_py_flashinfer_ragged_mha_prefill.py:208- 建议:给
_test_prefill_correctness增加v_size_per_head参数,让hidden_size_v、v.reshape、_create_kv_cache与参考实现按 vo 维取值,新增size_per_head=192 / v_size_per_head=128正确性用例(FP8 子类自动继承)。resolve_paged_backend/resolve_ragged_backend/window_left_from_sliding_window是不依赖 CUDA 的纯函数,直接 parametrize 覆盖("auto",192,128)->"fa2"、("auto",128,128)->"auto"、("fa3",192,128)->"fa3"与 window 的 0/1/负数边界。apply_attention_sink对照朴素实现(显式把exp(b_h)加入分母)比数值,用已知(out,lse)把 log2 域假设固化成断言。_scatter_write在对称维下与append_paged_kv_cache比对,再构造含NULL_BLOCK_IDX的 block table 断言被改道 token 只落到块 0。invokeSwaMhaPagedAttnPlan补 gtest 覆盖全 NULL / 部分 NULL / 无 NULL。
- 建议:给
- 新模型端到端测试标记 manual,本 PR 无任何可在 CI 回归的端到端覆盖 @
rtp_llm/test/model_test/BUILD:14- 建议:按仓库既有 smoke 约定补一个少层数用例(
smoke_test()宏 +*_4layers裁剪权重 + golden JSON):裁一份 4 层左右、同时含 GA 与 SWA 层、保留 QK=192/V=128 非对称头维的权重,用gpu_type限定机型并纳入对应架构 suite。test_mimo_v25.py可保留为需完整权重的人工验收用例。若短期内无法准备裁剪权重,请在 PR description 中明确说明该新模型无 CI 覆盖及后续补齐计划。
- 建议:按仓库既有 smoke 约定补一个少层数用例(
- 机器专属开发脚本携带个人机器路径提交到公开仓库根目录 @
run_mimo_v25_test.sh:41- 建议:倾向不把该脚本纳入本 PR:
--override_repository=accl_ep_rpm这类内部构建环境适配更适合放在内部脚本目录而非开源仓库根目录。若确需保留:1) 删除所有指向具体个人机器的默认路径,CHECKPOINT_PATH改为必填、缺失时明确报错;2) cache 搜索改为读bazel info output_base/repository_cache,或仅在ACCL_EP_RPM_PATH与当前用户自己的 cache 下查找,不遍历其它用户家目录;3)--config与TP_SIZE改为可覆盖的环境变量(--jobs已支持BAZEL_JOBS);4) 移入内部脚本目录并在测试文档中说明用途。
- 建议:倾向不把该脚本纳入本 PR:
Non-blocking Suggestions
P2
- SWA 的 NULL 页消毒只接到 device planner,另外三条 planner 入口仍会把负页号交给 FlashInfer @
rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:1365- 建议:二选一收口:(1) 给
fillParams与fillDecodeCudaGraphParams也加上is_sliding_window,让四条 planner 路径共享同一消毒契约;(2) 或在PyFlashinferDecodeAttnOp.prepare()的分支条件里补上or self.window_left >= 0,与PyFlashinferPrefillPagedAttnOp.prepare保持一致,并对确实只有 host 输入的窗口组显式报错。无论选哪种,建议同时让support_cuda_graph()在window_left >= 0时也返回 False,并在注释写明原因是 replay planner 未消毒(而非非对称),避免后续放宽该谓词时静默踩中。
- 建议:二选一收口:(1) 给
- fill_params_mha_device 的 pybind 签名变更未同步到类型存根 @
rtp_llm/models_py/bindings/cuda/FlashInferMlaParams.cc:750- 建议:重新生成或手工补齐
rtp_llm/ops/librtp_compute_ops/rtp_llm_ops.pyi中fill_params_mha_device的签名,加上is_sliding_window: bool = False,保持 pybind 与类型存根一致。
- 建议:重新生成或手工补齐
- 非对称 decode 在 enable_cuda_graph 打开时无可用实现,且失败信息不可诊断 @
rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:1533- 建议:在
PyFlashinferDecodeAttnOp.prepare()中让use_prefill_for_decode短路掉:1395起的 CUDA graph 接线分支(该路径本就不参与捕获),并在MiMoV25Model构造或prepare_fmha_impl处显式校验:非对称 K/V 暂不支持enable_cuda_graph,抛出指明「请关闭 enable_cuda_graph」的错误,而不是让用户看到can not find mha type。
- 建议:在
- rope cache 兜底路径缺长度校验,且构建失败被静默吞没退化为只用 base theta @
rtp_llm/models_py/modules/factory/attention/cuda_impl/base_rotary_embedding_op.py:76- 建议:把长度校验提到两条分支之外,对最终返回值统一断言
data.size(0) >= max_position_embeddings,不满足则抛出而非静默 memo。收窄异常捕获:仅对确实可退化的情形(如max_position_embeddings <= 0的 warmup,见 RopeCache.cc:133)返回 None,其余按raise ... from e上抛;若坚持保留兜底,至少加 warning 记录 rope_config 摘要与异常,并在_apply_rope的动态分支里对style != RopeStyle.Base或scale != 1.0显式拒绝。另外checkRopeCache(RopeCache.cc:171-174)只比较used/dim/base/defined,docstring 中「actually matches」的说法弱于实际,建议改为准确表述或在 Python 侧补齐比较。
- 建议:把长度校验提到两条分支之外,对最终返回值统一断言
- SWA plan kernel 的页写入循环无 max_blocks_per_bs 上界保护 @
rtp_llm/models_py/bindings/cuda/kernels/mha_paged_attn_plan.cu:96- 建议:在 kernel 内把循环上界改为
min(pages_self, max_blocks_per_bs),或在launchMhaPagedAttnPlan增加一条TORCH_CHECK,用 host 侧已知量校验ceil(max_seq_len / seq_size_per_block) <= max_blocks_per_bs,让不变量违背时 fail-fast 而不是静默越界。同时建议在mha_paged_attn_plan.h的契约注释里把这条宽度要求写成显式前置条件。
- 建议:在 kernel 内把循环上界改为
- SWA 组开启前缀复用后保留全长 KV,窗口收益消失且与注释、离线估算假设均矛盾 @
rtp_llm/cpp/cache/HybridPoolConfigCreator.cc:321- 建议:先确定取舍再统一两侧口径:若要保留窗口收益,为该 desc 提供独立于全局
linear_step的保留步长,或让active_tail_blocks在 reuse 打开时仍参与释放判定;若接受全长保留,则修正 mimo_v25.py:64-67 的注释,并把_eval_hybrid_kv_cache_mem_size的 SWA 分支改为按实际保留量计量。无论哪种,请在 PR 描述中给出开启前缀复用前后block_num/ 最大并发的实测对比,便于容量回滚决策。
- 建议:先确定取舍再统一两侧口径:若要保留窗口收益,为该 desc 提供独立于全局
- 混合注意力 KV 显存估算有三处计量口径问题 @
rtp_llm/config/model_config.py:278- 建议:复用
hybrid_kv_cache.py的三分口径显式区分 LINEAR / SLIDING_WINDOW / 全局(LINEAR 计 0 或按linear_attention_config状态尺寸计费;若本次不修请在 docstring 写明已知偏差);GA 分支改用swa_config.ga_kv_head_num or self.attn_config.kv_head_num,或在_parse_swa_config中断言二者相等;并与apply_layer_num_override保持一致校验len(pattern) == self.num_layers,不满足时抛ValueError或退回同构公式并打 warning,不要静默产出 0。
- 建议:复用
- local KV head 计算公式在 spec 与 pool config 两处重复实现,仅靠注释维持一致 @
rtp_llm/cpp/cache/HybridPoolConfigCreator.cc:73- 建议:收敛为单一入口,例如在
MHAKVCacheSpec上提供static uint32_t resolveLocalKvHeads(const KVCacheSpecDesc&, const AttentionConfigs&, const ParallelismConfig&),由build()与mhaLocalKvHeadNum()共同调用,以类型系统而非注释保证一致性。
- 建议:收敛为单一入口,例如在
- v_size_per_head 的哨兵语义在 C++/Python 多处各自重复解析,且权重切分处无回退 @
rtp_llm/cpp/model_utils/AttentionConfig.h:91- 建议:在
AttentionConfigs绑定处补只读入口暴露权威解析结果(如.def_property_readonly("v_size_per_head_resolved", &AttentionConfigs::vSizePerHead)),把上述四处 Python 回退改为读取该入口,使哨兵解析只保留 C++ 一份实现;裸字段保留可写以供 checkpoint 解析赋值。同时在get_sp_tensor_kv_asym入口对v_size_per_head > 0做显式断言,避免哨兵值穿透到张量切分逻辑。
- 建议:在
- pickle 契约变更未接入既有 config_pickle_test.py @
rtp_llm/cpp/pybind/ConfigInit.cc:1905- 建议:在
config_pickle_test.py按既有模式补三类用例:KVCacheSpecDesc22 元组往返(先给全部字段赋非默认值,断言kv_head_num_override/v_size_per_head_override及后移的group_type/reuse/capacity/memory/tail/cp逐一还原,尤其覆盖嵌套tail非空)、非法长度(20、21)被拒、以及CacheTailPolicyDesc的两条__setstate__路径。
- 建议:在
- CacheTailPolicyDesc pickle 的宽松长度校验与同文件兄弟实现不一致且实际不可达 @
rtp_llm/cpp/pybind/ConfigInit.cc:1818- 建议:统一策略二选一:若不需跨版本兼容(当前实际情况),改为与兄弟实现一致的严格
t.size() != 3并删除死分支;若确需兼容,则让KVCacheSpecDesc也接受 20 元组并在该分支把group_type..cp从t[14..19]读回。无论哪种,都应把c.validate_tail_blocks = t[1].cast<std::optional<bool>>();提到if之外消除重复赋值。
- 建议:统一策略二选一:若不需跨版本兼容(当前实际情况),改为与兄弟实现一致的严格
- 新增 SWA/非对称视图用例的覆盖缺口:fail-fast 分支、FP8 scale、FULL 组、FULL+SWA 混合拓扑均未覆盖 @
rtp_llm/cpp/cache/test/KVCacheLayoutViewTest.cc:179- 建议:本文件已有
EXPECT_ANY_THROW用法(:212-223),补齐成本很低:1) 增一个local_kv_heads=2且kernel_seq_size < physical_seq_size的非对称用例并EXPECT_ANY_THROW;2) 在SwaMhaUsesKernelView补EXPECT_FALSE(layer.k_cache.defined())与EXPECT_FALSE(layer.v_cache.defined());3) 把SwaMhaBuildsAsymmetricKvViews参数化为{SWA, FULL}×{无 scale, 有 scale};4) 新增一条同一CacheTopology内放 FULL(kernel_seq=physical)与 SWA(kernel_seq<physical)两个 MHA 组的用例,分别getLayerCache(0, tag)断言各自seq_size_per_block、shape、K/V 偏移独立正确。
- 建议:本文件已有
- /tmp 下固定名目录 + 符号链接跟随写入 @
run_mimo_v25_test.sh:71- 建议:改为无条件使用
mktemp -d(配合trap ... EXIT清理),或在创建前用[[ -L "$dir" ]]拒绝符号链接并用-o(属主)校验目录归属;标记文件不足以证明目录归属。
- 建议:改为无条件使用
- KV page 大小硬编码为 64,跨页不变量在 ROCm/PPU 上会静默失效 @
rtp_llm/test/model_test/test_mimo_v25.py:52- 建议:在
smoke_args中显式加上--seq_size_per_block {KV_PAGE_SIZE},把常量变成被固定的输入而不是对默认值的假设;或从服务端返回信息读取实际页大小并据此计算min_new_tokens。SLIDING_WINDOW建议在setUpClass里读取 checkpoint 目录下config.json的sliding_window并断言与常量一致,一旦 checkpoint 变化立刻显式失败。
- 建议:在
- test_prefix_cache_reuse 在 REUSE_CACHE 关闭时空跑通过,且冷启断言与重试机制冲突 @
rtp_llm/test/model_test/test_mimo_v25.py:410- 建议:
REUSE_CACHE关闭时改用self.skipTest("reuse_cache disabled")或@unittest.skipUnless(...)显式跳过,让跳过在测试报告里可见,而不是让断言消失。冷启探针请求改用retry_times=1,让重试导致的污染直接暴露为请求失败;或把断言放宽为assertGreater(second_reuse, first_reuse),只校验「重复请求复用变多」这个真正要验证的性质。
- 建议:
- 以模型语义输出作为硬断言,易产生无法 root-cause 的偶发失败 @
rtp_llm/test/model_test/test_mimo_v25.py:359- 建议:把「是否正确」的判定从模型语义改为可确定性校验的量:例如对同一 prompt 在
top_k=1下比对两次输出完全一致,或与一份 golden 输出比对(沿用 smoke 的 golden 机制)。若确实要保留 needle 召回,建议降级为告警日志或独立的「模型质量」用例,不与 KV 布局正确性混在同一断言里。
- 建议:把「是否正确」的判定从模型语义改为可确定性校验的量:例如对同一 prompt 在
- setUpClass 启动失败时不回收服务进程,泄漏 4 卡显存 @
rtp_llm/test/model_test/test_mimo_v25.py:176- 建议:将 raise 包在
try/finally或在抛异常前显式调用cls._server.stop_server();也可参考rtp_llm/test/prompt_scoring_test.py的做法,checkpoint 缺失时用unittest.SkipTest而非硬失败,减少这条路径被触发的频率。
- 建议:将 raise 包在
P3
- swaMatchBegin 中 window_tokens==1 的特例分支与通用公式完全等价 @
rtp_llm/cpp/cache/HybridKVCacheAllocator.cc:70- 建议:删除
window_tokens == 1分支,仅保留window_tokens == 0的兼容早退与通用公式;若担心读者误解「窗口为 1 时不复用任何块」,改为在通用公式上方补一行注释说明该退化结果。
- 建议:删除
- SWA 复用的越界写入被静默丢弃,窗口全失配又会清零 FULL 复用且无日志或指标 @
rtp_llm/cpp/cache/HybridKVCacheAllocator.cc:225- 建议:把 :225 的静默边界改为显式不变量断言(
RTP_LLM_CHECK_WITH_INFO(...)后无条件setAt),使不变量破坏在开发期即暴露。同时当has_tail_groups && min_full_reuse_blocks > 0 && reuse_blocks_len == 0时打印一次可聚合的 warning 或上报一个 counter,使「声明了窗口但复用全失效」可被观测。
- 建议:把 :225 的静默边界改为显式不变量断言(
- 复用回退循环对同一 cache key 重复 matchSingleKey,最坏开销随声明窗口线性放大 @
rtp_llm/cpp/cache/HybridKVCacheAllocator.cc:162- 建议:把已查询的
key_pos → block_idx结果缓存在循环外的std::unordered_map(或按key_pos下标的 vector)中跨 pos 复用;或改为先自顶向下一次性求出每个 key 的 SWA 命中情况,再用「最长连续命中后缀」直接推出可行pos,把复杂度降回 O(keys)。
- 建议:把已查询的
- to_string() 无条件输出全零 SWA 块,却仍不打印决定层型的 hybrid_attention_types @
rtp_llm/cpp/config/ConfigModules.cc:202- 建议:仅在
enable_hybrid_attention为真时输出 SWA 子块;同时补充 pattern 的紧凑摘要,例如各层型计数(ga_layers: N, swa_layers: M)加首个 SWA 层索引——48 层逐个枚举偏冗长,计数加索引即可满足排查需要。
- 建议:仅在
- 新增 SwaAttentionConfig 落在 C++ 配置层但全仓无 C++ 消费者 @
rtp_llm/cpp/config/ConfigModules.h:637- 建议:若短期内没有 C++ 侧读取计划,可评估将该结构放到
rtp_llm/config/py_config_modules.py(仓内既有的 Python-only 配置载体),省去 C++ 重编、pybind 注册顺序约束与.pyi存根维护成本。若确实为后续 C++ 缓存/kernel 路径预留,请在结构上方注明预期消费方,使「当前无 C++ 引用」成为有意为之而非遗漏。
- 建议:若短期内没有 C++ 侧读取计划,可评估将该结构放到
- AttentionConfigs 同时暴露两个 V 头维字段,add_sink_bias 在两个结构上双写 @
rtp_llm/cpp/pybind/ConfigInit.cc:1669- 建议:不必改字段名(会波及既有 MLA 接入代码),建议在绑定处或
.pyi上给两者加区分性说明(v_head_dim标注「MLA only」、v_size_per_head标注「MHA asymmetric V dim, 0 = same as size_per_head」),或在配置校验处加互斥检查(use_mla为真时禁止设v_size_per_head,反之禁止设v_head_dim)。add_sink_bias建议在AttentionConfigs侧改为由swa_attention_config派生的只读属性,或在 C++ 校验处断言二者一致;若确需双写,请注明哪一侧是权威来源。
- 建议:不必改字段名(会波及既有 MLA 接入代码),建议在绑定处或
- 量化包装分派用 literal getattr 读未声明的鸭子类型标记 @
rtp_llm/utils/model_weight.py:1787- 建议:把标记提升为
WeightModule(或相关基类)上声明的类属性,默认False,让四个子类的覆写受类型检查约束,is_mimo_v25_weight直接读属性而不用getattr默认值;或统一走注册表显式登记的方式,与is_v4_weight的名字前缀分派收敛为同一套机制。
- 建议:把标记提升为
- TP_SIZE 暴露了已知非法的可配置项,失败信息会误导 @
rtp_llm/test/model_test/test_mimo_v25.py:151- 建议:在
setUpClass读到tp_size后立即校验,不等于 4 时raise ValueError(或self.skipTest)并直接说明「FP8 checkpoint 将 fused QKV 存为 4 个原生 slab,只能以 tp_size=4 加载」;或者干脆去掉该环境变量、直接写常量 4,让约束在代码里就是硬的。
- 建议:在
Checklist Findings (20 fail / 67 total)
General Principles Checklist
- [6.1] Architecture — 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全 → issue
CacheTailPolicyDesc pickle 的宽松长度校验与同文件兄弟实现不一致且实际不可达
本次把__setstate__放宽为t.size() != 2 && t.size() != 3,但同文件所有兄弟 policy desc(CacheMemoryPolicyDesc:1800 严格!=1、CacheCpPolicyDesc:1844 严格!=5)以及外层KVCacheSpecDesc(:1905 严格!=22)全部使用精确长度校验。该宽松分支实际不可达:CacheTailPolicyDesc只作为KVCacheSpecDesc::tail出现,旧的 20 元组会先在 :1905 抛错,子级的 2 元组分支永不进入。此外 :1822 与 :1825 两个分支重复赋值同一句c.validate_tail_blocks。两处对同一次结构变更采取不同兼容策略,容易让维护者误判「这里支持跨版本 pickle」。 - [6.1] Architecture — 分层边界:新概念在正确层级,不泄漏内部 → issue
新增 SwaAttentionConfig 落在 C++ 配置层但全仓无 C++ 消费者
全仓检索swa_attention_config/SwaAttentionConfig:C++ 侧仅出现在ConfigModules.h:637,650、ConfigModules.cc:190-207(to_string)与ConfigInit.cc:904-930(绑定)三处,无任何引擎代码读取其字段。真实读写方全在 Python:models/mimo_v25.py:220写入,model_config.py:276、mimo_v25_weight.py:381、models_py/model_desc/mimo_v25.py:137,184读取。对比同层的LinearAttentionConfig确有 C++ 消费者,二者性质不同。考虑到ModelConfig本就是跨进程配置载体,此处按可辩护的取舍记录而非缺陷。 - [6.1] Architecture — 可观测性:日志/指标/超时可操作、非噪声 → issue
to_string() 无条件输出全零 SWA 块,却仍不打印决定层型的 hybrid_attention_types
HybridAttentionConfig::to_string()无条件拼接"swa_attention_config: {\n" << swa_attention_config.to_string() << "\n}\n"(:206-207),即使enable_hybrid_attention == false,绝大多数模型的启动配置 dump 会固定多出 5 行全零字段。同时hybrid_attention_types至今仍未输出——该字段是唯一决定「哪些层是滑窗层」的信息,窗口尺寸与两套 head 数是否生效完全取决于它。结果是 MiMo 能看到window_size: 128、swa_kv_head_num: 8,却无法回答哪些层实际走了 SWA 路径:有诊断价值的字段缺席,零值噪声常驻。 - [6.1] Architecture — 回滚路径:风险行为存在运维回滚手段 → issue
SWA 组开启前缀复用后保留全长 KV,窗口收益消失且与注释、离线估算假设均矛盾
createHybridAttentionPoolConfig硬编码config.linear_step = 1。在SWAKVCacheGroup中 step=1 使step_hit = ((i+1)%1)==0恒真:shouldAllocateBlock(SWAKVCacheGroup.cc:22-23)在 reuse 打开时对每个槽位都分配,removeSkippedBlocks(:229-231)对所有槽位continue、不释放任何块。而 MiMo 的 SWA desc 显式设enable_prefix_reuse=True(mimo_v25.py:80),故请求级占用等于全长,并非该 desc 注释(mimo_v25.py:64-67)所称「occupancy bounded by active_tail_blocks」。同时_eval_hybrid_kv_cache_mem_size(model_config.py:287-291)按min(window_size, max_seq_len)计量 SWA 层,两侧 - [6.1] Architecture — 状态不变量:创建/更新/失败/重试/回滚路径有效 → issue
AttentionConfigs 同时暴露两个 V 头维字段,add_sink_bias 在两个结构上双写
ConfigInit.cc:1669已有v_head_dim(MLA 分支的 V 头维,位于// mla config段),:1687新增的v_size_per_head是 MHA 非对称模型的 V 头维。两个名字在同一 Python 对象上并存,语义差异只体现在 C++ 头文件注释里,.pyi中两者均为裸int且无区分说明。同名可写字段add_sink_bias也同时暴露在SwaAttentionConfig(:918)与AttentionConfigs(:1688)上,一致性完全依赖models/mimo_v25.py的手工赋值,绑定层无任何不变量约束;后续新增模型若只设置其中一个,绑定层不会报错,只在数值比对阶段暴露。 - [6.1] Architecture — 错误语义:fail-fast/retry/fallback/silent 行为显式 → issue
TP_SIZE 暴露了已知非法的可配置项,失败信息会误导
模块 docstring 与setUpClass注释都写明「tp_size 4 is the only value the FP8 checkpoint loads at」,已核实rtp_llm/models/mimo_v25_weight.py:300确有assert tp == QKV_QUANT_SHARDS。但代码仍从环境变量读取int(os.environ.get("TP_SIZE", "4"))并原样透传给--tp_size/--world_size。任何非 4 的取值都会在权重加载深处断言失败,最终只表现为RuntimeError("MiMo V2.5 server failed to start. See test_mimo_v25.log."),使用者必须去翻服务端日志才能知道真正原因,等于提供了一个注定失败的旋钮。 - [6.1] Quality — 逻辑变更未混入无关格式化 → issue
机器专属开发脚本携带个人机器路径提交到公开仓库根目录
该脚本以new file mode ******新增在仓库根目录(与 WORKSPACE / .bazelrc 同级)。第 41 行 bazel cache 搜索列表写死了某位开发者个人数据盘下的 cache 目录,第 101 行CHECKPOINT_PATH默认值指向某位开发者家目录下的模型目录——两处都会随公开仓库对外发布,泄漏内部机器布局与用户名,且在其他机器上默认不可用。第 42 行/home/*/.cache/bazel还会遍历所有用户家目录做rpm2cpio探测。第 94 行--config=cuda12_9、第 98 行TP_SIZE=4亦写死。全仓搜索确认无任何 CI 配置、文档或脚本引用它。 - [6.1] Software Engineering — DIP:高层策略不依赖非必要具体细节 → issue
量化包装分派用 literal getattr 读未声明的鸭子类型标记
is_mimo_v25_weight用return bool(getattr(src_weight_info, "is_mimo_v25", False))判定,被per_block_fp8_quant_weight.py:291用于把 MiMo 权重从基类support()中排除。标记is_mimo_v25 = True分散声明在models/mimo_v25_weight.py:161,192,203,209四个类上,另有一处同样的 literalgetattr(mimo_v25_weight.py:234),没有基类字段或协议声明,任何一处漏写或拼错都不会被类型检查发现——只会在注册表匹配检查处暴露为一条与根因无关的错误。 - [6.1] Software Engineering — DRY:重复非平凡逻辑被抽取或显式复用 → issue
swaMatchBegin 中 window_tokens==1 的特例分支与通用公式完全等价
if (window_tokens == 1) { return end + 1; }(:70-72)。代入通用公式:history_tokens = 0→history_start = first_new_token = (end+1)*block_tokens→begin = (end+1)*block_tokens / block_tokens = end + 1,逐字相同。该分支既不改变行为也不规避除零(block_tokens已由std::max(1, ...)保底),只多一条读者需单独验证语义的路径。 - [6.1] Software Engineering — KISS/YAGNI:无投机性抽象 → issue
新增 SwaAttentionConfig 落在 C++ 配置层但全仓无 C++ 消费者
全仓检索swa_attention_config/SwaAttentionConfig:C++ 侧仅出现在ConfigModules.h:637,650、ConfigModules.cc:190-207(to_string)与ConfigInit.cc:904-930(绑定)三处,无任何引擎代码读取其字段。真实读写方全在 Python:models/mimo_v25.py:220写入,model_config.py:276、mimo_v25_weight.py:381、models_py/model_desc/mimo_v25.py:137,184读取。对比同层的LinearAttentionConfig确有 C++ 消费者,二者性质不同。考虑到ModelConfig本就是跨进程配置载体,此处按可辩护的取舍记录而非缺陷。 - [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue
KV page 大小硬编码为 64,跨页不变量在 ROCm/PPU 上会静默失效
第 52 行注释称KV_PAGE_SIZE = 64是--seq_size_per_block默认值,但rtp_llm/config/server_config_setup.py:414-429的默认值随平台而变:检测到/dev/kfd为 16、/dev/alixpu为 256,其余才是 64;而setUpClass组装的smoke_args(:159-168)并未传--seq_size_per_block。在 256 的平台上max_new_tokens=64根本不会跨越页边界,但第 340 行assertGreaterEqual(output_len, KV_PAGE_SIZE)依然通过——用例的核心保证失效却不报错。SLIDING_WINDOW = 128同理硬编码复制了 checkpoint 的sliding_window,若权重侧该值变大,assertGreater会空转通过。 - [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue
以模型语义输出作为硬断言,易产生无法 root-cause 的偶发失败
第 359-367 行要求生成文本中精确包含_FILLER下一行的前 5 个词,第 287-293 行要求原样召回 needle 码。这两条断言的成立取决于模型是否恰好延续 12 行填充轮转、以及是否原样复读 needle,而不取决于本 PR 的 KV 布局/注意力代码是否正确。一旦失败,失败信息把原因指向「全局注意力层或窗口掩码有问题」,但实际可能只是模型没这么续写,与仓库「每次 smoke 失败都必须定位根因、不得归为已知抖动」的要求相冲突。 - [6.1] Tests — 边界 case 覆盖(空、单元素、最大值) → issue
复用回退循环对同一 cache key 重复 matchSingleKey,最坏开销随声明窗口线性放大
外层for (; pos >= 0 && has_tail_groups; --pos)(:136)每次迭代都对 SWA 组重新扫描[begin, pos]全区间(:164-171),相邻迭代的查询区间高度重叠却不复用结果。改动前每个 pos 只做 1 次matchSingleKey,现在做window_blocks次,最坏总量为min_full_reuse_blocks × window_blocks。MiMo 当前 window=128 / page=64 仅约 2 倍,且未提供 benchmark,故按 P3 记录;但代码对窗口大小无上限约束,更大窗口下放大系数会同比上升。
RTP-LLM Checklist
- [G] 跨语言框架陷阱 — Java volatile/Atomic → issue
pickle 契约变更未接入既有 config_pickle_test.py
本次改动了两个 pickle 契约:KVCacheSpecDesc由 20 扩为 22 元组并重排t[14]…t[21](:1894-1928),CacheTailPolicyDesc由 2 扩为 3 元组(:1811-1828)。仓库已有专用测试rtp_llm/cpp/pybind/config_pickle_test.py,含 legacy/previous 元组兼容用例与test_fabricated_short_layouts_are_rejected的非法长度校验模式,但该文件只覆盖 GrammarConfig,不在本 PR 改动列表内。新增的 22 元组往返、下标重排与 tail 兼容分支全部无测试保护,下标错位类回归无法被 CI 捕获。 - [H] 测试与 CI — 新增、迁移或删除测试必须证明目标行为仍在 CI 中执行;DISABLED_、#if 0、open_skip、导入即失败、未被 target 引用的用例不得计为覆盖 → issue
以模型语义输出作为硬断言,易产生无法 root-cause 的偶发失败
第 359-367 行要求生成文本中精确包含_FILLER下一行的前 5 个词,第 287-293 行要求原样召回 needle 码。这两条断言的成立取决于模型是否恰好延续 12 行填充轮转、以及是否原样复读 needle,而不取决于本 PR 的 KV 布局/注意力代码是否正确。一旦失败,失败信息把原因指向「全局注意力层或窗口掩码有问题」,但实际可能只是模型没这么续写,与仓库「每次 smoke 失败都必须定位根因、不得归为已知抖动」的要求相冲突。 - [I] 代码质量 — 同一功能用统一工具函数 → issue
AttentionConfigs 同时暴露两个 V 头维字段,add_sink_bias 在两个结构上双写
ConfigInit.cc:1669已有v_head_dim(MLA 分支的 V 头维,位于// mla config段),:1687新增的v_size_per_head是 MHA 非对称模型的 V 头维。两个名字在同一 Python 对象上并存,语义差异只体现在 C++ 头文件注释里,.pyi中两者均为裸int且无区分说明。同名可写字段add_sink_bias也同时暴露在SwaAttentionConfig(:918)与AttentionConfigs(:1688)上,一致性完全依赖models/mimo_v25.py的手工赋值,绑定层无任何不变量约束;后续新增模型若只设置其中一个,绑定层不会报错,只在数值比对阶段暴露。
Python Static-First Checklist
- [P.A] 静态结构与类型纪律 — 禁止 getattr/setattr literal 访问 → issue
量化包装分派用 literal getattr 读未声明的鸭子类型标记
is_mimo_v25_weight用return bool(getattr(src_weight_info, "is_mimo_v25", False))判定,被per_block_fp8_quant_weight.py:291用于把 MiMo 权重从基类support()中排除。标记is_mimo_v25 = True分散声明在models/mimo_v25_weight.py:161,192,203,209四个类上,另有一处同样的 literalgetattr(mimo_v25_weight.py:234),没有基类字段或协议声明,任何一处漏写或拼错都不会被类型检查发现——只会在注册表匹配检查处暴露为一条与根因无关的错误。 - [P.B] 错误处理 — 异常链用 raise X from Y → issue
rope cache 兜底路径缺长度校验,且构建失败被静默吞没退化为只用 base theta
共享单例分支做了shared.data.size(0) >= max_position_embeddings校验(:71-73),但else的get_rope_cache(...)(:76)返回值直接进 memo。已核实getRopeCache(RopeCache.cc:98-106)在 Yarn 下把rope_config.max_pos而非入参传给genYarnCache,后者 cache 长度实为arange(max_pos * rope_scale)(:69);乘积小于请求长度时_apply_rope(:149-160)用pos_ids索引会越界读。同时except Exception: return None(:77-79)不记日志、不分类,兜底分支(:161-167)只传rope_config.base,丢弃 scale/mscale/yarn 修正——而getRopeCache:110-111对不支持的 style 正是抛异常。 - [P.B] 错误处理 — 禁止 bare except 或静默吞异常 → issue
rope cache 兜底路径缺长度校验,且构建失败被静默吞没退化为只用 base theta
共享单例分支做了shared.data.size(0) >= max_position_embeddings校验(:71-73),但else的get_rope_cache(...)(:76)返回值直接进 memo。已核实getRopeCache(RopeCache.cc:98-106)在 Yarn 下把rope_config.max_pos而非入参传给genYarnCache,后者 cache 长度实为arange(max_pos * rope_scale)(:69);乘积小于请求长度时_apply_rope(:149-160)用pos_ids索引会越界读。同时except Exception: return None(:77-79)不记日志、不分类,兜底分支(:161-167)只传rope_config.base,丢弃 scale/mscale/yarn 修正——而getRopeCache:110-111对不支持的 style 正是抛异常。 - [P.B] 错误处理 — 资源获取用 with context manager → issue
setUpClass 启动失败时不回收服务进程,泄漏 4 卡显存
第 176-180 行在start_server返回 False 时直接raise RuntimeError。unittest 在setUpClass抛异常时不会调用tearDownClass,因此第 187 行的stop_server()不会执行。此时cls._server仍是类属性、引用长期存活,MagaServerManager.__del__只可能在解释器退出时才触发。健康检查超时(如 backend 已加载权重但卡在 NCCL)属于最常见的失败形态,此路径下 start_server 子进程及其 backend/frontend 子进程会持续占用 4 张卡。
Strengths
- 契约注释把适用边界钉在 API 层而非调用点:
mha_paged_attn_plan.h:43-44明确写出「不要用于全注意力组,那里负页号是 bug,静默改写会掩盖它」,并用template<bool kSanitizeNullPages>双特化(mha_paged_attn_plan.cu:21,200,224)而非运行时分支,全注意力路径零额外判断。 - 多处把「静默算错」改成 fail-fast:
OpDefs.h:230拒绝对非对称 K/V 按 kernel-block 粒度切分(否则会返回一个看起来合理的转置视图);MemoryLayoutStrategy.cc:250拒绝对非对称 K/V 做kv/2分区;select_paged_kv_cache:263-267在非对称却拿到空k_cache时抛出并指明根因;mimo_v25.py:136-143对 ckpt 的swa_head_dim/attention_projection_layout做前置断言。 apply_attention_sink的推导(out * D/(D+exp(b)) = out * sigmoid(logD - b))与_LN2换底均标注了 FlashInfer 的具体源文件位置,并说明为何不走sinks=参数(fa2/fa3 分支静默丢弃)与 JIT-only wrapper;window_left_from_sliding_window也主动纠正了 HFsliding_window与 FlashInferwindow_left之间的 off-by-one 而非直接透传。_resolve_cos_sin_cache正确识别出getRopeCacheOnce的std::call_once语义(RopeCache.cc:142-167)会让 39 个 SWA 层(theta 1e4)误用 9 个 GA 层(theta 1e7)的 cache,修的是一处原本会静默算错的真实缺陷。- 向后兼容默认值选得稳:
prefix_reuse_window_tokens默认 0 时swaMatchBegin()直接返回end,与改动前逐字节等价;is_paged_group用spec->type == OpaqueState(HybridPoolConfigCreator.cc:253-255)排除 DSv4 的固定分配状态池,既有模型 paged 预算不变。 - 跨语言契约同步无裂缝:
SwaAttentionConfig在 C++ 结构体、pybind、.pyi三处字段名/类型一致,作为第 4 个成员追加在HybridAttentionConfig末尾使既有 3 参聚合构造仍有效,且ConfigInit.cc:904注明了「必须先于按引用暴露它的HybridAttentionConfig注册」这一顺序约束。 has_asymmetric_kv_head_dim(utils.py:8)刻意只读v_size_per_head这个纯形状事实,不读sliding_window/add_sink_bias,避免误改 DeepSeek-V4 的 impl 选路;trt.py:43与trtllm_gen.py:586/652/728四处support()一致接入同一谓词,未各自手写判断。KVCacheLayoutViewTest.cc:179-200不只断言形状,还断言k_cache/v_cache的strides()完全相等与v_cache.storage_offset()==4,正好锁住「FlashInfer paged kernel 只带一套 stride」这一硬约束;逐值复算{3,2,8,6}/{96,48,6,1}与 head_dim=4 / v_head_dim=2 的推导一致。
| // so one key covers cp_scale raw blocks. A zero window keeps the historical | ||
| // single-tail behavior used by generic SWA groups that do not declare a | ||
| // prefix-reuse window; MiMo supplies the explicit window in its descriptor. | ||
| int swaMatchBegin(const KVCacheGroup& group, int end, int cp_scale) { |
There was a problem hiding this comment.
[P1] SWA 多块前缀复用(swaMatchBegin + CP 槽位映射)无任何单测覆盖
新逻辑含两处易错算术:swaMatchBegin() 用 (end+1)*block_tokens - (window-1) 反推起始块(:74-78),以及非 compact CP 布局下 logical_pos = (canonical_pos+1)*cp_scale - 1(:223-224)。全仓 grep prefix_reuse_window_tokens / swaMatchBegin,仅命中 mimo_v25.py、KVCacheSpecDesc、CacheGroupType、CacheConfig、ConfigInit、pyi 与本文件,rtp_llm/cpp/cache/test/ 下 0 命中;本 PR 唯一的 cache 测试改动只覆盖张量视图。取错块不会崩溃,只会让 attention 读到窗口外 KV,或在窗口内留下 NULL 空洞(plan kernel 替入保留块 0),属静默错答。
建议: 在 HybridTypeKVCacheAllocatorTest.cc(已有 SWA policy 脚手架)补 reuseCache 用例:(1) window_tokens=0 与旧单 tail 行为逐位等价;(2) 窗口跨 2~3 个 block 时 [begin,pos] 全部落位、更早槽位保持 NULL;(3) 窗口内任一 key 未命中时 pos 逐级回退;(4) 边界值 window_tokens=1、window_tokens <= seq_size_per_block、begin 被 max(0,...) 截断至 0。在 HybridKVCacheAllocatorCPShardTest.cc 补 cp_size>1 下 compact 与非 compact 两种 logical_pos 映射断言。建议同时把 swaMatchBegin 暴露为可直接单测的纯函数入口。
| with_kv_cache_block_ids=False, | ||
| ) | ||
|
|
||
| def test_asymmetric_head_dims_pin_fa2_backend(self): |
There was a problem hiding this comment.
[P1] 非对称 K/V、attention sink 与窗口换算的数值正确性无 CI 覆盖,唯一新增用例不可证伪
新增四处纯计算逻辑失败形态均为「结果错但不报错」:apply_attention_sink(py_flashinfer_mha.py:212,依赖 FlashInfer lse 处于 log2 域这一内部细节)、window_left_from_sliding_window(:148,sliding_window-1)、_scatter_write(kv_cache_write_op.py:166)、invokeSwaMhaPagedAttnPlan(mha_paged_attn_plan.cu:213)。全仓检索这些符号仅命中实现与 .pyi,无测试引用。唯一新增用例只断言 attn_op.backend == "fa2"(:219),无对称配置反例——resolve_ragged_backend 若被改为无条件返回 fa2 仍通过;_test_prefill_correctness 在 :118/:137 全程用 size_per_head 表达 V,结构上无法表达非对称。
建议: 给 _test_prefill_correctness 增加 v_size_per_head 参数,让 hidden_size_v、v.reshape、_create_kv_cache 与参考实现按 vo 维取值,新增 size_per_head=192 / v_size_per_head=128 正确性用例(FP8 子类自动继承)。resolve_paged_backend / resolve_ragged_backend / window_left_from_sliding_window 是不依赖 CUDA 的纯函数,直接 parametrize 覆盖 ("auto",192,128)->"fa2"、("auto",128,128)->"auto"、("fa3",192,128)->"fa3" 与 window 的 0/1/负数边界。apply_attention_sink 对照朴素实现(显式把 exp(b_h) 加入分母)比数值,用已知 (out,lse) 把 log2 域假设固化成断言。_scatter_write 在对称维下与 append_paged_kv_cache 比对,再构造含 NULL_BLOCK_IDX 的 block table 断言被改道 token 只落到块 0。invokeSwaMhaPagedAttnPlan 补 gtest 覆盖全 NULL / 部分 NULL / 无 NULL。
| "//rtp_llm/test/utils:maga_server_manager", | ||
| ], | ||
| # Needs a real MiMo-V2.5 checkpoint (CHECKPOINT_PATH) and 4 GPUs. | ||
| tags = ["manual"], |
There was a problem hiding this comment.
[P1] 新模型端到端测试标记 manual,本 PR 无任何可在 CI 回归的端到端覆盖
tags = ["manual"] + timeout = "eternal",注释写明需真实 checkpoint 与 4 卡。manual 会让该 target 被 ... 通配排除,永不在 CI 执行。而本 PR 引入了新模型 mimo_v25、CacheGroupType::SWA 的 kernel-view 路径、非对称 K/V 的 k_cache/v_cache 视图、SWA plan kernel 与 backend 收敛逻辑,diff_paths 中没有任何 smoke 用例或 golden 数据变更;kv_cache_specs_test.py 已为 qwen3_next / kimi_linear 覆盖 spec-desc 构建,本次也未为 mimo_v25 补入。这些新路径的端到端行为因此完全没有自动化回归手段。
建议: 按仓库既有 smoke 约定补一个少层数用例(smoke_test() 宏 + *_4layers 裁剪权重 + golden JSON):裁一份 4 层左右、同时含 GA 与 SWA 层、保留 QK=192/V=128 非对称头维的权重,用 gpu_type 限定机型并纳入对应架构 suite。test_mimo_v25.py 可保留为需完整权重的人工验收用例。若短期内无法准备裁剪权重,请在 PR description 中明确说明该新模型无 CI 覆盖及后续补齐计划。
| local cache_root | ||
| for cache_root in \ | ||
| /root/.cache/bazel \ | ||
| /data1/renkun.ren/.cache/bazel \ |
There was a problem hiding this comment.
[P1] 机器专属开发脚本携带个人机器路径提交到公开仓库根目录
该脚本以 new file mode ****** 新增在仓库根目录(与 WORKSPACE / .bazelrc 同级)。第 41 行 bazel cache 搜索列表写死了某位开发者个人数据盘下的 cache 目录,第 101 行 CHECKPOINT_PATH 默认值指向某位开发者家目录下的模型目录——两处都会随公开仓库对外发布,泄漏内部机器布局与用户名,且在其他机器上默认不可用。第 42 行 /home/*/.cache/bazel 还会遍历所有用户家目录做 rpm2cpio 探测。第 94 行 --config=cuda12_9、第 98 行 TP_SIZE=4 亦写死。全仓搜索确认无任何 CI 配置、文档或脚本引用它。
建议: 倾向不把该脚本纳入本 PR:--override_repository=accl_ep_rpm 这类内部构建环境适配更适合放在内部脚本目录而非开源仓库根目录。若确需保留:1) 删除所有指向具体个人机器的默认路径,CHECKPOINT_PATH 改为必填、缺失时明确报错;2) cache 搜索改为读 bazel info output_base / repository_cache,或仅在 ACCL_EP_RPM_PATH 与当前用户自己的 cache 下查找,不遍历其它用户家目录;3) --config 与 TP_SIZE 改为可覆盖的环境变量(--jobs 已支持 BAZEL_JOBS);4) 移入内部脚本目录并在测试文档中说明用途。
Checklist: [6.1] 逻辑变更未混入无关格式化
| # their stale capacity sizes (MIN_CACHE_BATCH_SIZE), which corrupts | ||
| # plan's batch size. Route tensor-core through the host fill. | ||
| if attn_inputs.input_lengths.is_cuda and not self.use_tensor_core: | ||
| if self.use_prefill_for_decode or ( |
There was a problem hiding this comment.
[P2] SWA 的 NULL 页消毒只接到 device planner,另外三条 planner 入口仍会把负页号交给 FlashInfer
is_sliding_window 只存在于 fillParamsMhaDevice(FlashInferMlaParams.cc:626、:666),host 侧 fillParams 无对应变体。三条入口缺消毒:(1) decode prepare() 的分支条件是 use_prefill_for_decode or (input_lengths.is_cuda and not use_tensor_core)(:1365),缺少 paged prefill 那侧特意加的 or self.window_left >= 0(对比 :350),窗口层若 K/V 对称且走 tensor-core 就落到 else 的 fill_params(:1386);(2) prepare_for_cuda_graph_replay host 分支调 fill_params(:1426);(3) 其 device 分支调 fill_decode_cuda_graph_params(:1448),该接口无此形参。当前唯一阻断是 `support_cuda_...
建议: 二选一收口:(1) 给 fillParams 与 fillDecodeCudaGraphParams 也加上 is_sliding_window,让四条 planner 路径共享同一消毒契约;(2) 或在 PyFlashinferDecodeAttnOp.prepare() 的分支条件里补上 or self.window_left >= 0,与 PyFlashinferPrefillPagedAttnOp.prepare 保持一致,并对确实只有 host 输入的窗口组显式报错。无论选哪种,建议同时让 support_cuda_graph() 在 window_left >= 0 时也返回 False,并在注释写明原因是 replay planner 未消毒(而非非对称),避免后续放宽该谓词时静默踩中。
| pybind11::arg("forbid_realloc") = false, | ||
| "MHA-only device-resident planner — fills decode_page_indptr_d / paged_kv_last_page_len_d / page_indice_d via a single CUDA kernel, leaving MLA-only fields untouched") | ||
| pybind11::arg("forbid_realloc") = false, | ||
| pybind11::arg("is_sliding_window") = false, |
There was a problem hiding this comment.
[P2] fill_params_mha_device 的 pybind 签名变更未同步到类型存根
pybind 在 FlashInferMlaParams.cc:750 新增了 pybind11::arg("is_sliding_window") = false,但 rtp_llm/ops/librtp_compute_ops/rtp_llm_ops.pyi:26 仍是旧签名(末参为 forbid_realloc: bool = False),且该 .pyi 不在本 PR 的改动文件列表中。调用点 py_flashinfer_mha.py:363 与 :1380 均以关键字形式传 is_sliding_window=,在仓库声明的 pyright strict 下会报 unexpected keyword argument。同一 PR 已更新 libth_transformer_config.pyi(含补齐历史遗漏的 AttentionConfigs.sliding_window),说明存根文件是在维护的,此处属遗漏。
建议: 重新生成或手工补齐 rtp_llm/ops/librtp_compute_ops/rtp_llm_ops.pyi 中 fill_params_mha_device 的签名,加上 is_sliding_window: bool = False,保持 pybind 与类型存根一致。
| return True | ||
| # The asymmetric path runs RoPE and the KV write from Python, which the capture | ||
| # does not cover. | ||
| return not self.asymmetric_kv |
There was a problem hiding this comment.
[P2] 非对称 decode 在 enable_cuda_graph 打开时无可用实现,且失败信息不可诊断
support_cuda_graph() 对 asymmetric_kv 返回 False(:1533)。DECODE_MHA_IMPS 中其余候选均排除 192 head dim:xqa 要求 size_per_head in [64,128,256],trtllm_gen.py:652/728 本 PR 新增 has_asymmetric_kv_head_dim 直接 return False。attn_factory.py:202 在 is_cuda_graph 下拒绝后循环耗尽,抛出 Exception("can not find mha type")(:214),用户看不出真正原因。另外工厂是先构造实例(构造期即调 prepare())再检查 support_cuda_graph(),构造期仍会进入 :1395 的 graph 接线分支。
建议: 在 PyFlashinferDecodeAttnOp.prepare() 中让 use_prefill_for_decode 短路掉 :1395 起的 CUDA graph 接线分支(该路径本就不参与捕获),并在 MiMoV25Model 构造或 prepare_fmha_impl 处显式校验:非对称 K/V 暂不支持 enable_cuda_graph,抛出指明「请关闭 enable_cuda_graph」的错误,而不是让用户看到 can not find mha type。
| ): | ||
| data = shared.data | ||
| else: | ||
| data = get_rope_cache(rope_config, max_position_embeddings, interleave) |
There was a problem hiding this comment.
[P2] rope cache 兜底路径缺长度校验,且构建失败被静默吞没退化为只用 base theta
共享单例分支做了 shared.data.size(0) >= max_position_embeddings 校验(:71-73),但 else 的 get_rope_cache(...)(:76)返回值直接进 memo。已核实 getRopeCache(RopeCache.cc:98-106)在 Yarn 下把 rope_config.max_pos 而非入参传给 genYarnCache,后者 cache 长度实为 arange(max_pos * rope_scale)(:69);乘积小于请求长度时 _apply_rope(:149-160)用 pos_ids 索引会越界读。同时 except Exception: return None(:77-79)不记日志、不分类,兜底分支(:161-167)只传 rope_config.base,丢弃 scale/mscale/yarn 修正——而 getRopeCache:110-111 对不支持的 style 正是抛异常。
建议: 把长度校验提到两条分支之外,对最终返回值统一断言 data.size(0) >= max_position_embeddings,不满足则抛出而非静默 memo。收窄异常捕获:仅对确实可退化的情形(如 max_position_embeddings <= 0 的 warmup,见 RopeCache.cc:133)返回 None,其余按 raise ... from e 上抛;若坚持保留兜底,至少加 warning 记录 rope_config 摘要与异常,并在 _apply_rope 的动态分支里对 style != RopeStyle.Base 或 scale != 1.0 显式拒绝。另外 checkRopeCache(RopeCache.cc:171-174)只比较 used/dim/base/defined,docstring 中「actually matches」的说法弱于实际,建议改为准确表述或在 Python 侧补齐比较。
Checklist: [P.B] 异常链用 raise X from Y;[P.B] 禁止 bare except 或静默吞异常
| @@ -88,15 +94,15 @@ __global__ void mhaPagedAttnPlanKernel(const int32_t* __restrict__ input_lengths | |||
| if (kv_cache_block_id) { | |||
| const int32_t* row = kv_cache_block_id + tid * max_blocks_per_bs; | |||
| for (int j = 0; j < pages_self; ++j) { | |||
There was a problem hiding this comment.
[P2] SWA plan kernel 的页写入循环无 max_blocks_per_bs 上界保护
for (int j = 0; j < pages_self; ++j) { page_indice[p_start + j] = ...; }(:96-99)中 pages_self = ceil(seq_len / seq_size_per_block)(:56)来自序列长度,而 row = kv_cache_block_id + tid * max_blocks_per_bs(:95)的行宽来自 block table。循环没有 j < max_blocks_per_bs 约束,launchMhaPagedAttnPlan 也只校验 page_indice.numel() >= batch_size * max_blocks_per_bs(:162-164)。一旦某请求 pages_self > max_blocks_per_bs,读会跨到下一行、写会越过上界,构成设备端越界写。该循环早于本 PR 且当前不变量成立,但 SWA 变体新引入的「block table 全长绝对寻址、仅尾部物化」约定(mha_paged_attn_plan.h:36)...
建议: 在 kernel 内把循环上界改为 min(pages_self, max_blocks_per_bs),或在 launchMhaPagedAttnPlan 增加一条 TORCH_CHECK,用 host 侧已知量校验 ceil(max_seq_len / seq_size_per_block) <= max_blocks_per_bs,让不变量违背时 fail-fast 而不是静默越界。同时建议在 mha_paged_attn_plan.h 的契约注释里把这条宽度要求写成显式前置条件。
| } | ||
|
|
||
| // HybridAttentionConfig | ||
| std::string HybridAttentionConfig::to_string() const { |
There was a problem hiding this comment.
[P3] to_string() 无条件输出全零 SWA 块,却仍不打印决定层型的 hybrid_attention_types
HybridAttentionConfig::to_string() 无条件拼接 "swa_attention_config: {\n" << swa_attention_config.to_string() << "\n}\n"(:206-207),即使 enable_hybrid_attention == false,绝大多数模型的启动配置 dump 会固定多出 5 行全零字段。同时 hybrid_attention_types 至今仍未输出——该字段是唯一决定「哪些层是滑窗层」的信息,窗口尺寸与两套 head 数是否生效完全取决于它。结果是 MiMo 能看到 window_size: 128、swa_kv_head_num: 8,却无法回答哪些层实际走了 SWA 路径:有诊断价值的字段缺席,零值噪声常驻。
建议: 仅在 enable_hybrid_attention 为真时输出 SWA 子块;同时补充 pattern 的紧凑摘要,例如各层型计数(ga_layers: N, swa_layers: M)加首个 SWA 层索引——48 层逐个枚举偏冗长,计数加索引即可满足排查需要。
Checklist: [6.1] 可观测性:日志/指标/超时可操作、非噪声
| SLIDING_WINDOW, | ||
| }; | ||
|
|
||
| struct SwaAttentionConfig { |
There was a problem hiding this comment.
[P3] 新增 SwaAttentionConfig 落在 C++ 配置层但全仓无 C++ 消费者
全仓检索 swa_attention_config / SwaAttentionConfig:C++ 侧仅出现在 ConfigModules.h:637,650、ConfigModules.cc:190-207(to_string)与 ConfigInit.cc:904-930(绑定)三处,无任何引擎代码读取其字段。真实读写方全在 Python:models/mimo_v25.py:220 写入,model_config.py:276、mimo_v25_weight.py:381、models_py/model_desc/mimo_v25.py:137,184 读取。对比同层的 LinearAttentionConfig 确有 C++ 消费者,二者性质不同。考虑到 ModelConfig 本就是跨进程配置载体,此处按可辩护的取舍记录而非缺陷。
建议: 若短期内没有 C++ 侧读取计划,可评估将该结构放到 rtp_llm/config/py_config_modules.py(仓内既有的 Python-only 配置载体),省去 C++ 重编、pybind 注册顺序约束与 .pyi 存根维护成本。若确实为后续 C++ 缓存/kernel 路径预留,请在结构上方注明预期消费方,使「当前无 C++ 引用」成为有意为之而非遗漏。
Checklist: [6.1] 分层边界:新概念在正确层级,不泄漏内部;[6.1] KISS/YAGNI:无投机性抽象
| @@ -1665,7 +1682,10 @@ PYBIND11_MODULE(libth_transformer_config, m) { | |||
| .def_readwrite("o_groups", &AttentionConfigs::o_groups) | |||
There was a problem hiding this comment.
📍 实际位置 rtp_llm/cpp/pybind/ConfigInit.cc:1669(不在 diff 展示范围内,就近挂载)
[P3] AttentionConfigs 同时暴露两个 V 头维字段,add_sink_bias 在两个结构上双写
ConfigInit.cc:1669 已有 v_head_dim(MLA 分支的 V 头维,位于 // mla config 段),:1687 新增的 v_size_per_head 是 MHA 非对称模型的 V 头维。两个名字在同一 Python 对象上并存,语义差异只体现在 C++ 头文件注释里,.pyi 中两者均为裸 int 且无区分说明。同名可写字段 add_sink_bias 也同时暴露在 SwaAttentionConfig(:918)与 AttentionConfigs(:1688)上,一致性完全依赖 models/mimo_v25.py 的手工赋值,绑定层无任何不变量约束;后续新增模型若只设置其中一个,绑定层不会报错,只在数值比对阶段暴露。
建议: 不必改字段名(会波及既有 MLA 接入代码),建议在绑定处或 .pyi 上给两者加区分性说明(v_head_dim 标注「MLA only」、v_size_per_head 标注「MHA asymmetric V dim, 0 = same as size_per_head」),或在配置校验处加互斥检查(use_mla 为真时禁止设 v_size_per_head,反之禁止设 v_head_dim)。add_sink_bias 建议在 AttentionConfigs 侧改为由 swa_attention_config 派生的只读属性,或在 C++ 校验处断言二者一致;若确需双写,请注明哪一侧是权威来源。
Checklist: [6.1] 状态不变量:创建/更新/失败/重试/回滚路径有效;[I] 同一功能用统一工具函数
| (class attribute ``is_mimo_v25 = True``). MiMo uses the standard W slot names, so | ||
| quant-wrapping dispatch discriminates by instance marker instead of the name prefix | ||
| DSv4 relies on (`is_v4_weight`).""" | ||
| return bool(getattr(src_weight_info, "is_mimo_v25", False)) |
There was a problem hiding this comment.
[P3] 量化包装分派用 literal getattr 读未声明的鸭子类型标记
is_mimo_v25_weight 用 return bool(getattr(src_weight_info, "is_mimo_v25", False)) 判定,被 per_block_fp8_quant_weight.py:291 用于把 MiMo 权重从基类 support() 中排除。标记 is_mimo_v25 = True 分散声明在 models/mimo_v25_weight.py:161,192,203,209 四个类上,另有一处同样的 literal getattr(mimo_v25_weight.py:234),没有基类字段或协议声明,任何一处漏写或拼错都不会被类型检查发现——只会在注册表匹配检查处暴露为一条与根因无关的错误。
建议: 把标记提升为 WeightModule(或相关基类)上声明的类属性,默认 False,让四个子类的覆写受类型检查约束,is_mimo_v25_weight 直接读属性而不用 getattr 默认值;或统一走注册表显式登记的方式,与 is_v4_weight 的名字前缀分派收敛为同一套机制。
Checklist: [6.1] DIP:高层策略不依赖非必要具体细节;[P.A] 禁止 getattr/setattr literal 访问
|
|
||
| # tp_size 4 is the only value the FP8 checkpoint loads at, see module docstring. | ||
| # max_seq_len must hold the long-context prompt below plus its output. | ||
| tp_size = int(os.environ.get("TP_SIZE", "4")) |
There was a problem hiding this comment.
[P3] TP_SIZE 暴露了已知非法的可配置项,失败信息会误导
模块 docstring 与 setUpClass 注释都写明「tp_size 4 is the only value the FP8 checkpoint loads at」,已核实 rtp_llm/models/mimo_v25_weight.py:300 确有 assert tp == QKV_QUANT_SHARDS。但代码仍从环境变量读取 int(os.environ.get("TP_SIZE", "4")) 并原样透传给 --tp_size / --world_size。任何非 4 的取值都会在权重加载深处断言失败,最终只表现为 RuntimeError("MiMo V2.5 server failed to start. See test_mimo_v25.log."),使用者必须去翻服务端日志才能知道真正原因,等于提供了一个注定失败的旋钮。
建议: 在 setUpClass 读到 tp_size 后立即校验,不等于 4 时 raise ValueError(或 self.skipTest)并直接说明「FP8 checkpoint 将 fused QKV 存为 4 个原生 slab,只能以 tp_size=4 加载」;或者干脆去掉该环境变量、直接写常量 4,让约束在代码里就是硬的。
Checklist: [6.1] 错误语义:fail-fast/retry/fallback/silent 行为显式
Add MiMo-V2.5 support on top of the current
mainbranch.MiMo-V2.5 is a hybrid-attention model with 9 global-attention layers and 39 sliding-window-attention layers. The two attention types use different KV-cache geometries and require independent cache groups.
Changes
ga_kvandswa_kvcache pools