执行摘要
本PR将WanVideo模型在AMD平台上的RoPE（旋转位置编码）实现从Triton内核替换为aiter HIP内核，通过减少内存分配和形状转换开销，实现内核级约40%的性能提升，端到端去噪时间减少约0.5%。变更已通过准确性测试验证，风险较低，是AMD平台多模态模型性能优化的典型案例。
功能与动机
为什么做？ 原Triton实现rope_cached_thd_positions_2c_fwd每次调用需分配两个新输出张量，并需进行reshape/view_as操作，导致性能开销。aiter HIP内核rope_cached_2c_fwd_inplace能以单内核完成k和v的RoPE计算，且支持原位操作，从而优化内存使用和计算效率。
解决什么问题？ 提升WanVideo模型在AMD平台上的推理性能，减少去噪步骤的延迟。PR body中提供了详细的基准测试数据：内核级时间从2,922.2 μs降至1,334.6 μs，端到端去噪时间从115.14秒降至114.68秒。
实现拆解
实现涉及两个文件的关键改动：
-
配置层 (python/sglang/multimodal_gen/runtime/models/utils.py):
python
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
新增_use_aiter标志，基于环境变量和HIP平台检测动态启用aiter优化。
-
核心逻辑层 (python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py):
- 在
forward方法中，当_use_aiter为真时，替换原有_apply_rotary_emb调用。
- 关键步骤：
python
q_sbhd = query.view(num_tokens, 1, query_shape[-2], query_shape[-1])
k_sbhd = key.view(num_tokens, 1, key_shape[-2], key_shape[-1])
rope_cached_2c_fwd_inplace(q_sbhd, k_sbhd, cos_sbhd, sin_sbhd, ...)
将张量重塑为sbhd格式后调用HIP内核，再恢复原始形状。
评论区精华
Review中唯一但关键的讨论围绕导入安全性展开：
HaiShaw: "Please not to do conditionless aiter import, do within _is_hip."
Jackycheng0808: "Remove the unconditional aiter import; import it only when use_aiter is enabled..."
核心洞察：初始提交使用try:无条件导入aiter模块，可能在非HIP环境引发错误。通过将导入移至_use_aiter条件内，确保平台兼容性。这体现了在平台特异性优化中维护代码健壮性的最佳实践。
风险与影响
风险：
- 兼容性：依赖aiter库和HIP平台，若环境配置不当可能回退到原有路径，需确保回退逻辑正确。
- 正确性：尽管准确性测试通过（CLIP均值0.9899、SSIM均值0.8657等指标达标），但内核替换需确保边界情况处理一致。
- 性能收益：端到端改进仅0.5%，可能受其他瓶颈限制，实际收益需生产验证。
影响：
- 用户：AMD平台用户获得轻微性能提升，功能不变。
- 系统：减少内存分配和内核调用，潜在降低系统开销。
- 团队：延续AMD优化脉络，与近期PR（如#22424的LayerNorm优化）形成技术协同。
关联脉络
本PR是AMD平台性能优化系列的一部分：
- #22424：使用aiter CK layernorm2d减少NSA索引器内核启动，与本PR同为内核替换优化。
- #22329：为AMD MORI-EP新增环境变量，与本PR共享通过环境变量控制平台特性的模式。
- #22089：为Qwen3-ASR添加流式语音识别，反映多模态领域的并行投入。
演进趋势：sglang项目正持续深化平台特异性优化（尤其是AMD），通过内核级替换（Triton→aiter/HIP）提升性能，同时注重通过环境变量和条件导入保持代码可维护性。
参与讨论