执行摘要

该 PR 修复了 AITER 上游 API 重命名导致 vLLM 在 ROCm 平台下 MLA dual RMS norm 融合运行时崩溃的问题。通过引入旧新 API 自动探测和 hot path 性能优化，实现了前后兼容，并在 pass manager 中增加守卫条件，避免编译异常。

功能与动机

AITER PR#2958 将 fused_qk_rmsnorm 函数重命名为私有的 _fused_qk_rmsnorm，且要求调用方预先分配输出缓冲区（q_out, k_out）。这使得依赖于该函数的 MLADualRMSNormFusionPass 在加载新 AITER 时抛出 ImportError。该 PR 旨在恢复此融合功能在新旧 AITER 版本上的可用性。

实现拆解

1. 更新兼容性检测函数 check_aiter_fused_qk_rmsnorm（vllm/_aiter_ops.py）：优先尝试导入 _fused_qk_rmsnorm，失败后回退到 fused_qk_rmsnorm。结果缓存至模块级变量，避免重复检测。
2. 重构实现函数 _fused_mla_dual_rms_norm_impl（vllm/_aiter_ops.py）：只导入模块一次，使用 hasattr 判断新（_fused_qk_rmsnorm）旧（fused_qk_rmsnorm）API，分别调用并准备正确参数。消除了异常开销，并防止无关 ImportError 被吞没。
3. Pass Manager 守卫（vllm/compilation/passes/pass_manager.py）：引入 check_aiter_fused_qk_rmsnorm() 条件，仅当函数可用时才注册 MLADualRMSNormFusionPass，否则静默跳过。
4. 测试验证：使用 AITER 0.1.13（旧 API）和新版 AITER 运行 tests/compile/passes/test_fuse_mla_dual_rms_norm.py，以及在 Kimi-K2 模型上完成端到端推理确认。

vllm/_aiter_ops.py

核心修复文件：适配 AITER API 重命名，重构检查与实现函数以同时支持新旧 API

def check_aiter_fused_qk_rmsnorm() -> bool:
    '''检查 AITER 是否提供 fused_qk_rmsnorm。支持新旧 API。'''
    global _AITER_HAS_FUSED_QK_RMSNORM
    if _AITER_HAS_FUSED_QK_RMSNORM is None:
        try:
            from aiter.ops.fused_qk_norm_rope_cache_quant import _fused_qk_rmsnorm
            _AITER_HAS_FUSED_QK_RMSNORM = True
        except (ImportError, ModuleNotFoundError, AttributeError):
            try:
                from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_rmsnorm
                _AITER_HAS_FUSED_QK_RMSNORM = True
            except (ImportError, ModuleNotFoundError, AttributeError):
                _AITER_HAS_FUSED_QK_RMSNORM = False
    return _AITER_HAS_FUSED_QK_RMSNORM
​
def _fused_mla_dual_rms_norm_impl(x1, x1_weight, x2, x2_weight, x1_epsilon, x2_epsilon):
    '''使用 import-once + hasattr 分发。'''
    try:
        import aiter.ops.fused_qk_norm_rope_cache_quant as aiter_ops
    except ImportError as exc:
        raise ImportError('fused_qk_rmsnorm requires AITer >= PR #2442.') from exc
    if hasattr(aiter_ops, '_fused_qk_rmsnorm'):
        return aiter_ops._fused_qk_rmsnorm(q_out=None, q=x1, q_weight=x1_weight, q_eps=x1_epsilon, k_out=None, k=x2, k_weight=x2_weight, k_eps=x2_epsilon)
    if hasattr(aiter_ops, 'fused_qk_rmsnorm'):
        return aiter_ops.fused_qk_rmsnorm(q=x1, q_weight=x1_weight, q_eps=x1_epsilon, k=x2, k_weight=x2_weight, k_eps=x2_epsilon)
    raise ImportError('fused_qk_rmsnorm requires AITer >= PR #2442.')

评论区精华

• gemini-code-assist[bot] 指出原实现的重复 import 和异常捕获在热路径上存在性能隐患，建议改用 import-once + hasattr。该建议被作者采纳，重构了 _fused_mla_dual_rms_norm_impl。
• AndreasKaratzas 要求添加 TODO 并链接至上游 AITER issue（#3207）以跟踪 API 稳定化。作者按要求添加了注释和链接。

风险与影响

• 风险：低。AITER 未来可能进一步变更 API，但 PR 已提供的前后兼容模式可降低硬失败；遗留的旧路径可依赖 TODO 跟踪清理。
• 影响：仅限 ROCm 平台使用 AITER MLA fused dual RMS norm 的用户（如 Kimi-K2 模型）。对 CUDA 等其他平台无影响。

关联脉络

• 关联历史 PR #42409（同一文件 vllm/_aiter_ops.py，调整 AITER fused AR RMSNorm 门控条件），说明 vllm/_aiter_ops.py 是 AITER 功能在 vLLM 中的主要适配层，随 AITER 演进持续更新。
• 上游 issue #3207 跟踪 AITER 侧 API 稳定化，待 resolve 后 vLLM 可移除兼容分支。