执行摘要
本PR为AMD MORI-EP新增环境变量SGLANG_MORI_PREALLOC_MAX_RECV_TOKENS，允许用户配置接收令牌的预分配数量，以在内存占用和缓冲区溢出风险间取得平衡。同时引入向后兼容性检查函数，并更新相关文档。变更主要影响AMD平台的MoE调度模块，为内存敏感场景提供调优能力。
功能与动机
根据PR body，主要动机是：
- 添加环境变量：
SGLANG_MORI_PREALLOC_MAX_RECV_TOKENS，允许用户自定义MORI-EP的令牌预分配数量。
- 提供向后兼容性：通过
check_mori_compatibility函数处理不同MORI库版本的参数兼容性。
- 文档修复：更新环境变量文档和服务器参数帮助文本。
这解决了用户需要灵活控制内存占用与性能平衡的需求，特别是在内存受限的AMD硬件环境中。
实现拆解
实现涉及三个文件：
| 文件 |
关键变更 |
说明 |
python/sglang/srt/layers/moe/token_dispatcher/moriep.py |
新增check_mori_compatibility函数；在init_mori_op中读取环境变量并传递给MORI配置。 |
核心逻辑，通过动态反射检查MORI库兼容性，并集成新环境变量。 |
docs/references/environment_variables.md |
新增环境变量条目：SGLANG_MORI_PREALLOC_MAX_RECV_TOKENS。 |
提供使用指南：有效范围1到world_size*SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK，默认0表示最大值，设置过小可能导致缓冲区溢出。 |
python/sglang/srt/server_args.py |
修正帮助文本，将“DeepEP MoE”扩展为“DeepEP or MoriEP MoE”。 |
次要更新，反映更广泛的MoE支持。 |
关键代码片段：
def check_mori_compatibility(kwargs: dict) -> None:
"""Remove kwargs not accepted by the installed mori's EpDispatchCombineConfig."""
import dataclasses
config_cls = mori.ops.EpDispatchCombineConfig
valid_kwargs = {f.name for f in dataclasses.fields(config_cls)}
invalid_kwargs = set(kwargs.keys()) - valid_kwargs
for arg in invalid_kwargs:
logger.warning(f"[MORI compat] Removing incompatible argument {arg} ")
del kwargs[arg]
评论区精华
review讨论较少，但包含关键点：
@billishyahao Can you leave comments in code/doc on how SGLANG_MORI_PREALLOC_MAX_RECV_TOKENS should be used and default?
作者通过提交补充了文档描述，明确了环境变量的作用、默认值和风险警告，解决了reviewer的疑虑。
风险与影响
风险：
- 兼容性风险：
check_mori_compatibility依赖运行时反射，可能在不同MORI版本中行为不一致。
- 配置风险：环境变量设置过小可能导致缓冲区溢出，文档已警告但缺乏运行时验证。
- 回归风险：修改
init_mori_op参数传递逻辑，可能影响现有功能，但UT测试通过提供了基本保障。
影响：
- 用户：AMD平台用户获得内存调优能力，需自行权衡内存占用与溢出风险。
- 系统：扩展了MoE调度配置灵活性，增强平台适应性。
- 团队：需关注MORI库版本演进，确保兼容性逻辑稳定。
关联脉络
与近期PR的关联：
- PR 22424：同属AMD平台优化，涉及性能调优和内核调度。
- PR 20089：同属MoE分发路径（EP）相关功能，本PR扩展了MORI-EP支持。
- PR 22335：同属AMD平台问题修复，涉及兼容性处理。
这反映了sglang仓库在AMD平台上持续优化MoE和调度能力的趋势，特别是内存管理和性能调优方向。
参与讨论