Skip to content

mimo v2.5 support - #1351

Open
RK-PHG wants to merge 1 commit into
alibaba:mainfrom
RK-PHG:mimov25_support
Open

mimo v2.5 support#1351
RK-PHG wants to merge 1 commit into
alibaba:mainfrom
RK-PHG:mimov25_support

Conversation

@RK-PHG

@RK-PHG RK-PHG commented Aug 29, 2026

Copy link
Copy Markdown

Add MiMo-V2.5 support on top of the current main branch.

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

  • Add MiMo-V2.5 model registration and model description.
  • Add MiMo-V2.5 checkpoint and FP8 weight loading.
  • Add hybrid attention configuration for:
    • 9 global-attention layers
    • 39 sliding-window-attention layers
    • 128-token sliding window
    • distinct GA/SWA KV head configurations
  • Extend KV-cache descriptors and layout handling to support:
    • per-group KV head count overrides
    • asymmetric K/V head dimensions
    • separate ga_kv and swa_kv cache pools
    • hybrid KV-cache memory estimation
  • Implement window-aware, block-granular SWA prefix-cache matching.
  • Update attention, rotary embedding, KV-cache write, FlashInfer, and paged-attention integration paths required by MiMo-V2.5.
  • Add the MiMo-V2.5 test launcher with ACCL repository configuration and required runtime environment.

@LLLLKKKK LLLLKKKK left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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=1window_tokens <= seq_size_per_blockbeginmax(0,...) 截断至 0。在 HybridKVCacheAllocatorCPShardTest.cccp_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_vv.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 覆盖及后续补齐计划。
  • 机器专属开发脚本携带个人机器路径提交到公开仓库根目录 @ 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) --configTP_SIZE 改为可覆盖的环境变量(--jobs 已支持 BAZEL_JOBS);4) 移入内部脚本目录并在测试文档中说明用途。

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) 给 fillParamsfillDecodeCudaGraphParams 也加上 is_sliding_window,让四条 planner 路径共享同一消毒契约;(2) 或在 PyFlashinferDecodeAttnOp.prepare() 的分支条件里补上 or self.window_left >= 0,与 PyFlashinferPrefillPagedAttnOp.prepare 保持一致,并对确实只有 host 输入的窗口组显式报错。无论选哪种,建议同时让 support_cuda_graph()window_left >= 0 时也返回 False,并在注释写明原因是 replay planner 未消毒(而非非对称),避免后续放宽该谓词时静默踩中。
  • fill_params_mha_device 的 pybind 签名变更未同步到类型存根 @ rtp_llm/models_py/bindings/cuda/FlashInferMlaParams.cc:750
    • 建议:重新生成或手工补齐 rtp_llm/ops/librtp_compute_ops/rtp_llm_ops.pyifill_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.Basescale != 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 的契约注释里把这条宽度要求写成显式前置条件。
  • 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 / 最大并发的实测对比,便于容量回滚决策。
  • 混合注意力 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 按既有模式补三类用例:KVCacheSpecDesc 22 元组往返(先给全部字段赋非默认值,断言 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..cpt[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=2kernel_seq_size < physical_seq_size 的非对称用例并 EXPECT_ANY_THROW;2) 在 SwaMhaUsesKernelViewEXPECT_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_tokensSLIDING_WINDOW 建议在 setUpClass 里读取 checkpoint 目录下 config.jsonsliding_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 布局正确性混在同一断言里。
  • 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 而非硬失败,减少这条路径被触发的频率。

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,使「声明了窗口但复用全失效」可被观测。
  • 复用回退循环对同一 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++ 引用」成为有意为之而非遗漏。
  • 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++ 校验处断言二者一致;若确需双写,请注明哪一侧是权威来源。
  • 量化包装分派用 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 严格 !=1CacheCpPolicyDesc: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,650ConfigModules.cc:190-207to_string)与 ConfigInit.cc:904-930(绑定)三处,无任何引擎代码读取其字段。真实读写方全在 Python:models/mimo_v25.py:220 写入,model_config.py:276mimo_v25_weight.py:381models_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: 128swa_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_weightreturn 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 getattrmimo_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 = 0history_start = first_new_token = (end+1)*block_tokensbegin = (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,650ConfigModules.cc:190-207to_string)与 ConfigInit.cc:904-930(绑定)三处,无任何引擎代码读取其字段。真实读写方全在 Python:models/mimo_v25.py:220 写入,model_config.py:276mimo_v25_weight.py:381models_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_weightreturn 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 getattrmimo_v25_weight.py:234),没有基类字段或协议声明,任何一处漏写或拼错都不会被类型检查发现——只会在注册表匹配检查处暴露为一条与根因无关的错误。
  • [P.B] 错误处理 — 异常链用 raise X from Y → issue rope cache 兜底路径缺长度校验,且构建失败被静默吞没退化为只用 base theta
    共享单例分支做了 shared.data.size(0) >= max_position_embeddings 校验(:71-73),但 elseget_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),但 elseget_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 也主动纠正了 HF sliding_window 与 FlashInfer window_left 之间的 off-by-one 而非直接透传。
  • _resolve_cos_sin_cache 正确识别出 getRopeCacheOncestd::call_once 语义(RopeCache.cc:142-167)会让 39 个 SWA 层(theta 1e4)误用 9 个 GA 层(theta 1e7)的 cache,修的是一处原本会静默算错的真实缺陷。
  • 向后兼容默认值选得稳:prefix_reuse_window_tokens 默认 0 时 swaMatchBegin() 直接返回 end,与改动前逐字节等价;is_paged_groupspec->type == OpaqueStateHybridPoolConfigCreator.cc:253-255)排除 DSv4 的固定分配状态池,既有模型 paged 预算不变。
  • 跨语言契约同步无裂缝:SwaAttentionConfig 在 C++ 结构体、pybind、.pyi 三处字段名/类型一致,作为第 4 个成员追加HybridAttentionConfig 末尾使既有 3 参聚合构造仍有效,且 ConfigInit.cc:904 注明了「必须先于按引用暴露它的 HybridAttentionConfig 注册」这一顺序约束。
  • has_asymmetric_kv_head_dimutils.py:8)刻意只读 v_size_per_head 这个纯形状事实,不读 sliding_window / add_sink_bias,避免误改 DeepSeek-V4 的 impl 选路;trt.py:43trtllm_gen.py:586/652/728 四处 support() 一致接入同一谓词,未各自手写判断。
  • KVCacheLayoutViewTest.cc:179-200 不只断言形状,还断言 k_cache/v_cachestrides() 完全相等与 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) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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=1window_tokens <= seq_size_per_blockbeginmax(0,...) 截断至 0。在 HybridKVCacheAllocatorCPShardTest.cccp_size>1 下 compact 与非 compact 两种 logical_pos 映射断言。建议同时把 swaMatchBegin 暴露为可直接单测的纯函数入口。

with_kv_cache_block_ids=False,
)

def test_asymmetric_head_dims_pin_fa2_backend(self):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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_vv.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"],

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 覆盖及后续补齐计划。

Comment thread run_mimo_v25_test.sh
local cache_root
for cache_root in \
/root/.cache/bazel \
/data1/renkun.ren/.cache/bazel \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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) --configTP_SIZE 改为可覆盖的环境变量(--jobs 已支持 BAZEL_JOBS);4) 移入内部脚本目录并在测试文档中说明用途。

Checklist: [6.1] 逻辑变更未混入无关格式化

@LLLLKKKK LLLLKKKK left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Code Review - PR #1351 (non-blocking suggestions)

25 条 P2/P3 建议,不阻塞合并。阻塞判定与完整摘要见上一条 review。

# 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 (

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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) 给 fillParamsfillDecodeCudaGraphParams 也加上 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,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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.pyifill_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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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:202is_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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] rope cache 兜底路径缺长度校验,且构建失败被静默吞没退化为只用 base theta

共享单例分支做了 shared.data.size(0) >= max_position_embeddings 校验(:71-73),但 elseget_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.Basescale != 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) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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: 128swa_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 {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] 新增 SwaAttentionConfig 落在 C++ 配置层但全仓无 C++ 消费者

全仓检索 swa_attention_config / SwaAttentionConfig:C++ 侧仅出现在 ConfigModules.h:637,650ConfigModules.cc:190-207to_string)与 ConfigInit.cc:904-930(绑定)三处,无任何引擎代码读取其字段。真实读写方全在 Python:models/mimo_v25.py:220 写入,model_config.py:276mimo_v25_weight.py:381models_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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📍 实际位置 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))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] 量化包装分派用 literal getattr 读未声明的鸭子类型标记

is_mimo_v25_weightreturn 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 getattrmimo_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"))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[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 行为显式

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants