# PR #42311 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Model] [Perf] Use flatten for Qwen3.5's GDN output projection
- 合并时间：2026-05-18 16:14
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/42311

---

# 执行摘要

- 一句话：用 flatten 替换 einops rearrange 提升 GDN 输出投影性能
- 推荐动作：值得合并，变更微小但明确有效，且经过充分的性能与精度验证。建议精读的审阅者关注：`flatten(-2)` 与 `rearrange` 的语义等价性确认，以及 eager 模式下加速比的量化证据。该 PR 展示了消除无关 Python 开销的典型优化模式。

# 功能与动机

Qwen3.5 的 GDN 输出投影中，通过 `einops.rearrange(core_attn_out, "... h d -> ... (h d)")` 合并最后两个维度，每次前向调用都会触发昂贵的 Python 级字符串解析和验证。使用 PyTorch 等效的 `core_attn_out.flatten(-2)` 可避免该开销，从而提升性能。

# 实现拆解

1. 定位 `vllm/model_executor/layers/mamba/gdn_linear_attn.py` 文件中所有使用 `rearrange(core_attn_out, "... h d -> ... (h d)")` 的调用点，共有 3 处：`_output_projection` 方法、`forward_xpu` 方法和 `forward_cpu` 方法。
2. 将每个调用替换为 PyTorch 原生的 `core_attn_out.flatten(-2)`，并添加注释 `# ... h d -> ... (h d)` 以保持语义清晰。
3. 通过性能剖析验证：CUDA eager 模式下每调用平均从 23.261 us 降至 2.828 us（加速 8.23x），CPU 模式下从 41.741 us 降至 1.949 us（加速 21.41x）。
4. 通过困惑度测试验证精度：`einops.rearrange` 版本 PPL 为 19.5098，`flatten` 版本 PPL 为 19.5100，相对差异仅 0.0001%，与 HuggingFace 基准一致。
5. 无需修改测试配置，仅涉及源码层 3 行等量替换，无测试、配置或部署配套改动。

关键文件：
- `vllm/model_executor/layers/mamba/gdn_linear_attn.py`（模块 GDN 层；类别 source；类型 data-contract；符号 _output_projection, forward_xpu, forward_cpu）: 核心变更文件，在 3 个方法中将 einops.rearrange 替换为 torch.flatten(-2)。

关键符号：_output_projection, forward_xpu, forward_cpu

## 关键源码片段

### `vllm/model_executor/layers/mamba/gdn_linear_attn.py`

核心变更文件，在 3 个方法中将 einops.rearrange 替换为 torch.flatten(-2)。

```python
def _output_projection(
    self,
    core_attn_out: torch.Tensor,
    z: torch.Tensor,
    output: torch.Tensor,
    num_tokens: int,
):
    """Part 3: RMSNormGated + output linear projection."""
    z_shape_og = z.shape
    core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
    z = z.reshape(-1, z.shape[-1])
    core_attn_out = self.norm(core_attn_out, z)
    core_attn_out = core_attn_out.reshape(z_shape_og)
    # 使用 PyTorch flatten 替代 einops rearrange，
    # 避免字符串解析开销，加速比 CUDA 8.23x / CPU 21.41x
    core_attn_out = core_attn_out.flatten(-2)  # ... h d -> ... (h d)
    output[:num_tokens], _ = self.out_proj(core_attn_out)

```

# 评论区精华

主要讨论围绕变更的实际收益场景展开：
- 审阅者 ZJY0516 指出：解码阶段默认使用 CUDA Graph，不会触发 Python 级字符串解析，因此 eager 模式下的优化在 CUDA Graph 场景下可能无收益。
- 作者 rishaps 回应：虽然 CUDA Graph 会捕获内核，但 eager 模式（如预填充、CPU 推理、调试场景）中每调用 8-21x 的加速仍然有价值，且该变更降低了 CPU 推理延迟。
- 审阅者最终认可了变更，并批准合并。

- Eager 模式与 CUDA Graph 场景下的收益分歧 (performance): 认可 eager 场景下的收益，PR 被批准合并。

# 风险与影响

- 风险：风险极低：
 - 功能等价性：`flatten(-2)` 与 `rearrange("... h d -> ... (h d)")` 在语义上完全等价（合并最后两维），且已通过困惑度测试验证 PPL 无差异。
 - 兼容性：无新增依赖，仅移除对 `einops` 的一次调用，不影响其他用户。
 - 性能：仅在 CUDA Graph 捕获后无额外收益，但在 eager 模式下有显著提升。
- 影响：影响范围：
 - 直接影响 Qwen3.5 模型 GDN 层的所有前向路径（CUDA/XPU/CPU），覆盖 _output_projection、forward_xpu 和 forward_cpu。
 - 用户可见收益：在 eager 模式下预填充和 CPU 推理场景延迟降低；CUDA Graph 场景无变化。
 - 团队维护：减少对第三方库 `einops` 的依赖，代码更简洁且易于理解。
影响程度：低频调用路径，单次加速绝对值小（微秒级），但在大规模推理集群中可累积为可观的端到端延迟改善。

- 风险标记：低风险变更 , 性能微优化

# 关联脉络

- PR #42849 [Perf] Add do_not_specialize in fused FP8 RoPE kernel: 同为 vllm 项目中针对模型推理性能的微优化，属于同类性能改进模式。
- PR #42497 [Perf] Wire silu_and_mul_per_block_quant into TritonFP8MoE (MiniMax-M2): 同为模型层级的算子级性能优化，展示项目对减少 Python 层开销的持续关注。