PR body 中说明从 NVIDIA TensorRT-LLM 项目（PR #12163）移植内核，目的是优化内存访问并重用 SGLang 的自定义 all reduce v2 框架，以提升 MiniMax 模型在张量并行下的性能。引用性能结果：解码性能从 150 tps 提升到 157 tps。

该 PR 值得精读，特别是对于关注性能优化、JIT 内核设计和分布式计算的工程师。建议关注以下设计决策：

• eps 正确性处理的实现细节，确保数值稳定性。
• 自定义 all reduce v2 框架的扩展方式，如何支持新内核的块数配置。
• 模型集成中的环境变量使用和潜在回退机制，以平衡性能与鲁棒性。
  阅读时结合单元测试和基准脚本，以全面理解性能提升和风险点。

实现方案按模块拆解：

1. JIT 内核编译模块（python/sglang/jit_kernel/all_reduce.py）：新增 _jit_fused_parallel_qknorm_module、get_fused_parallel_qknorm_max_occupancy 和 fused_parallel_qknorm 函数，动态编译并调用 CUDA 内核。
2. CUDA 内核实现（python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh）：新增融合的 TP QK 归一化内核，重用自定义 all reduce push 缓冲区进行跨 GPU 通信。
3. 模型集成层（python/sglang/srt/models/minimax_m2.py）：添加 MiniMaxM2QKRMSNorm 类和 fused_tp_qknorm 函数，通过环境变量 SGLANG_USE_FUSED_PARALLEL_QKNORM 控制是否使用融合路径。
4. 分布式通信框架扩展（python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py）：扩展 CustomAllReduceV2 类以支持 max_pull_blocks 和 max_push_blocks 参数，优化内核调度。
5. 测试与基准：新增单元测试 test_tp_qknorm.py 和基准脚本 bench_tp_qknorm.py，确保正确性和性能验证。

技术风险具体如下：

• 正确性风险：eps 处理方式若未正确缩放，可能导致数值不稳定，影响模型输出精度。尽管讨论中认为已修正，但需依赖单元测试覆盖。
• 性能风险：推缓冲区大小固定为 1 MB（对应最多 131072 个令牌），在大批次场景下可能触发 Push buffer is too small 错误，导致回退或失败，缺乏优雅降级机制。
• 兼容性风险：新增环境变量 SGLANG_USE_FUSED_PARALLEL_QKNORM 未在文档中说明，用户可能不知如何启用优化，影响用户体验。
• 回归风险：修改了自定义 all reduce 框架的核心文件（如 common.cuh 和 custom_all_reduce.cuh），可能影响其他依赖该框架的功能，需确保测试充分。

影响范围和程度：

• 用户影响：MiniMax M2 模型用户可通过环境变量启用融合优化，解码性能提升约 4.7%，但需注意缓冲区限制和大批次场景。影响范围限于该模型用户，程度中等。
• 系统影响：新增 JIT 内核和测试代码，增加了系统复杂性，但提升了张量并行下的归一化效率，有利于整体性能优化。对系统核心路径有局部影响。
• 团队影响：开发团队需维护新的 CUDA 内核和集成逻辑，代码结构清晰且有测试支持，但可能增加维护负担。为后续类似优化提供了参考模板。
• 长期影响：推动 JIT 内核在分布式计算中的应用，可能促进更多模型性能优化，技术方向值得关注。