执行摘要
该PR修复了MiniMax M2.5模型在TP=16配置下因KV头复制导致的RMSNorm权重分片错误，解决了重复输出问题。通过重构MiniMaxM2RMSNormTP类，使其感知头复制并正确分片权重，同时修复了前向传播中的方差归约逻辑。这是一个针对特定模型和TP配置的关键bugfix，影响范围有限但修复了高端硬件下的模型可用性问题。
功能与动机
MiniMax M2.5模型使用8个KV头，当TP=16时，TP大小超过KV头数，多个TP rank必须共享同一个KV头。原有的MiniMaxM2RMSNormTP计算权重分片大小时使用hidden_size / attn_tp_size，对于K norm得到8d / 16 = 0.5d——一个非整数大小，导致错误的权重分片和归一化，最终产生重复或乱码输出。根本原因是MiniMaxM2RMSNormTP未感知到头级结构，假设权重维度总能被tp_size整除，这在tp_size > num_kv_heads时失效。
实现拆解
所有修改集中在python/sglang/srt/models/minimax_m2.py文件的MiniMaxM2RMSNormTP类中：
| 修改点 |
关键代码逻辑 |
目的 |
| 初始化 |
新增num_heads参数，根据QKVParallelLinear模式计算num_heads、num_head_replicas和head_dim |
确保权重大小始终为整数，支持头复制场景 |
| 权重加载器 |
从@staticmethod改为实例方法，使用attn_tp_rank // num_head_replicas计算分片索引 |
使复制rank正确加载相同权重分片 |
| 防御性断言 |
在__init__中添加整除性检查，在weight_loader中添加边界检查 |
提前捕获配置错误，防止静默失败 |
| 前向传播 |
修复方差归约逻辑，确保正确all-reduce和除法 |
解决reviewer指出的计算错误 |
评论区精华
reviewer JustinTong0323指出了三个关键问题，并推动了PR的改进：
初始化整除性验证："Both branches use integer division (//) without asserting exactness. If tp_size % num_heads != 0 (or vice versa), the remainder is silently dropped, causing incorrect shard calculations downstream."
方差归约逻辑错误："All-reduce sums all 16 ranks' variances (from 8 different heads × 2 replicas). Dividing by 2 gives sum_of_8_head_variances, but the correct per-layer variance requires dividing by 8..."
weight_loader边界检查："If the divisibility invariant from __init__ is violated, this will silently load incorrect weights — the model runs but produces wrong outputs with no error raised."
PR作者采纳了这些建议，添加了assert和边界检查，并修复了方差归约逻辑，体现了良好的代码审查文化。
风险与影响
技术风险：
- 回归风险：修改了核心的RMSNorm实现，如果头复制逻辑有误，可能导致其他TP配置下的模型输出错误。
- 兼容性风险：
MiniMaxM2RMSNormTP的__init__签名改变，破坏了向后兼容性，但仅影响MiniMax M2.5模型内部使用。
影响评估：
- 对用户：修复了TP=16时MiniMax M2.5模型的重复输出问题，提升了模型在高端硬件配置下的可用性。
- 对系统：仅影响MiniMax M2.5模型的RMSNorm实现，不涉及其他模型或子系统。
- 对团队：提供了头复制场景下权重分片的参考实现，可作为其他类似TP问题的解决模板。
关联脉络
从近期历史PR看，该PR与以下PR有相似之处：
- PR #22312：修复GDN内核以支持非连续张量输入，解决Qwen3.5-27B准确性回归问题，同为bugfix且涉及模型准确性。
- PR #22423：修复Flux.2模型TI2I准确性，通过对齐编码器、VAE和图像预处理行为，同为准确性修复。
这些PR共同反映了团队对模型准确性的持续关注，特别是在复杂配置（如TP>1、多模态）下的边缘case处理。该PR的解决方案——借鉴QKVParallelLinear的成熟模式——展示了代码复用的价值，为未来类似问题提供了参考。
参与讨论