执行摘要

• 一句话：修复GPT-J风格MRoPE硬编码并优化ROCm性能
• 推荐动作：此 PR 应合并。建议后续增加更多使用 adjacent-pair MRoPE 的模型测试，并考虑自动检测旋转风格（基于模型配置）以减少手动设置。Triton 内核的 mask 修改值得关注，是通用的正确性修复。

功能与动机

PR body 指出要 "Honor the existing adjacent-pair MRoPE setting instead of hardcoding NeoX pairing"。GLM-4.1V 等模型需要 GPT-J 风格 MRoPE，同时 ROCm 路径存在性能优化空间。

实现拆解

1. 在 _triton_mrope_forward 内核函数中添加 is_neox_style 编译时常量参数，控制旋转配对方式是 NeoX（前半/后半）还是 GPT-J（相邻对）。
2. 修改 mask 逻辑：在 is_interleaved 分支中增加 valid_mask 检查，确保 cos/sin 加载在有效范围内，避免越界。
3. 重写数据加载部分：根据 is_neox_style 选择加载左半部分（NeoX）还是相邻元素对（GPT-J），并在 ROCm 上使用连续的加载存储与寄存器-only 的 split/interleave 优化。
4. 使用 current_platform 抽象替代硬编码的 ROCm 检测，选择特定的 wave 启动形状，提高可移植性。
5. 在测试框架中为 MRoPETestInfo 添加 is_neox_style 字段，对 GLM-4.1V 模型指定 False，并在测试函数中传递该参数，确保正确验证。

关键文件：

• vllm/model_executor/layers/rotary_embedding/mrope.py（模块 位置编码；类别 source；类型 core-logic；符号 _triton_mrope_forward）: 核心 MRoPE 实现，修改 Triton 内核以支持两种旋转风格，并优化 ROCm 内存访问
• tests/kernels/core/test_mrope.py（模块 测试；类别 test；类型 test-coverage；符号 MRoPETestInfo, test_mrope, test_mrope_torch_compile_tracing）: 测试配套，验证两种 MRoPE 风格的正确性

关键符号：_triton_mrope_forward

关键源码片段

vllm/model_executor/layers/rotary_embedding/mrope.py

核心 MRoPE 实现，修改 Triton 内核以支持两种旋转风格，并优化 ROCm 内存访问

@triton.jit
def _triton_mrope_forward(
    q_ptr, k_ptr, cos, sin, num_tokens,
    n_qh: tl.constexpr, n_kh: tl.constexpr, hd: tl.constexpr, rd: tl.constexpr,
    pad_n_qh: tl.constexpr, pad_n_kh: tl.constexpr, pad_rd: tl.constexpr,
    mrope_section_t: tl.constexpr, mrope_section_h: tl.constexpr, mrope_section_w: tl.constexpr,
    is_interleaved: tl.constexpr, is_neox_style: tl.constexpr,
):
    # 前面地址计算省略 ...
    half_rd = rd // 2
    cos_offsets = tl.arange(0, pad_rd // 2)
​
    # 根据 interleaved 模式计算 mask
    if is_interleaved:
        valid_mask = cos_offsets < half_rd
        h_mask = valid_mask & ((cos_offsets % 3) == 1) & (cos_offsets <= 3 * mrope_section_h)
        w_mask = valid_mask & ((cos_offsets % 3) == 2) & (cos_offsets <= 3 * mrope_section_w)
        t_mask = valid_mask & ~(h_mask | w_mask)
    else:
        t_end = mrope_section_t
        h_end = t_end + mrope_section_h
        t_mask = cos_offsets < mrope_section_t
        h_mask = (t_end <= cos_offsets) & (cos_offsets < h_end)
        w_mask = (h_end <= cos_offsets) & (cos_offsets < half_rd)
​
    # 加载 cos/sin 三个模态的分量并求和
    t_cos_row = tl.load(t_cos + cos_offsets, mask=t_mask, other=0)
    h_cos_row = tl.load(h_cos + cos_offsets, mask=h_mask, other=0)
    w_cos_row = tl.load(w_cos + cos_offsets, mask=w_mask, other=0)
    cos_row = t_cos_row + h_cos_row + w_cos_row
    # sin 类似 ...
​
    # 根据 is_neox_style 选择旋转配对方式
    if is_neox_style:
        # NeoX：前一半和后一半配对
        rotary_offsets = tl.arange(0, pad_rd // 2)
        q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :]
        k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :]
        q_mask = (tl.arange(0, pad_n_qh)[:, None] < n_qh) & (rotary_offsets[None, :] < rd // 2)
        k_mask = (tl.arange(0, pad_n_kh)[:, None] < n_kh) & (rotary_offsets[None, :] < rd // 2)
        q_left = tl.load(q_ptr + q_offsets, mask=q_mask, other=0)
        q_right = tl.load(q_ptr + q_offsets + rd // 2, mask=q_mask, other=0)
        # 应用旋转并交错左右部分
    else:
        # GPT-J：相邻值配对，无需交错
        rotary_offsets = tl.arange(0, pad_rd)
        q_offsets = tl.arange(0, pad_n_qh)[:, None] * hd + rotary_offsets[None, :]
        k_offsets = tl.arange(0, pad_n_kh)[:, None] * hd + rotary_offsets[None, :]
        # 加载连续两个元素并应用旋转

tests/kernels/core/test_mrope.py

测试配套，验证两种 MRoPE 风格的正确性

class MRoPETestInfo(NamedTuple):
    model_name: str
    is_neox_style: bool = True # 默认 NeoX 风格
    atol: float = 1e-2
    rtol: float = 1.6e-2
    marks: list[pytest.MarkDecorator] = []
​
MODELS_TO_TEST = [
    MRoPETestInfo(
        model_name="zai-org/GLM-4.1V-9B-Thinking",
        is_neox_style=False, # GLM 使用 adjacent-pair MRoPE
    ),
    MRoPETestInfo(model_name="Qwen/Qwen2-VL-7B-Instruct"),
    MRoPETestInfo(model_name="Qwen/Qwen2-VL-72B-Instruct"),
    MRoPETestInfo(model_name="Qwen/Qwen2.5-VL-72B-Instruct"),
    MRoPETestInfo(model_name="Qwen/Qwen3-VL-4B-Instruct"),
    MRoPETestInfo(model_name="Qwen/Qwen3-VL-30B-A3B-Instruct"),
]
​
def test_mrope(default_vllm_config, model_name, model_info, tp_size, dtype, num_tokens):
    is_neox_style = model_info.is_neox_style # 从模型配置读取
    mrope_helper_class = get_rope(
        head_size=head_dim,
        max_position=max_position,
        is_neox_style=is_neox_style, # 传递风格参数
        rope_parameters=config.rope_parameters,
        dtype=dtype,
    )
    # ... forward_native 与 forward_cuda 比较

评论区精华

tjtanaa 评论指出：“This PR also changes the behaviour on NVIDIA, it is better to have someone that is able to validate the changes on NVIDIA to make sure there are no change in the performance.” 审核者 mgoin 随后回复“Validated on B300”，确认 NVIDIA B300 上性能无退化。

• NVIDIA 性能影响 (performance): mgoin 在 B300 上验证通过

风险与影响

• 风险：1. 核心旋转逻辑变更：所有使用 MRoPE 的多模态模型（Qwen2-VL、GLM-4.1V 等）均受此影响。默认 is_neox_style=True 保持原行为，但若模型配置错误（例如手动覆盖为 False 但实际使用 NeoX）会导致精度下降。
  2. 跨平台风险：NVIDIA 路径同样修改，虽经 B300 验证，但其他架构（如 A100）未测试。
  3. ROCm 优化风险：连续加载存储和 wave 形状优化对 ROCm 特定版本敏感，CI 覆盖有限。
  4. 测试覆盖：目前仅 GLM-4.1V 配置了 GPT-J 风格，其他使用 adjacent-pair 的模型未覆盖（如可能需要更多模型）。

• 影响：影响范围：中等。MRoPE 是 vLLM 中多模态模型（Qwen2-VL、GLM-4.1V 等）位置编码的核心组件。本次变更修复了 GPT-J 风格模型的正确性错误，对 NeoX 风格模型应无行为变化。
  性能影响：ROCm 路径获得优化（连续内存访问、寄存器-only 操作），NVIDIA 路径保持不变。
  兼容性：默认行为未变，但模型开发者可能需要明确指定 is_neox_style 在 rope_parameters 中。

• 风险标记：核心旋转逻辑变更, 影响多模态模型, ROCm 优化可能影响 NVIDIA, 测试覆盖不足（仅 GLM-4.1V 配置 GPT-J）

关联脉络

• 暂无明显关联 PR