作者：昇腾实战派

背景概述

在大模型推理场景中，Prefill阶段的吞吐性能直接影响整体服务时延，尤其在长输入、短输出的典型业务场景下，Prefill耗时成为TTFT（Time to First Token）的关键瓶颈。本文针对Atlas 800I A2平台上的模型Prefill性能问题展开分析，该模型在vLLM-Ascend v0.11.0环境下运行，时延更大。通过Profiling诊断与多维度优化，成功定位并解决快慢卡问题，显著提升Prefill阶段性能。

环境版本信息

问题定位：快慢卡现象与性能瓶颈分析

1. WithStack场景：CPU抢占导致的快慢卡

在开启调用栈追踪（withstack）模式下，Profiling结果显示存在明显的快慢卡现象，部分卡执行时延显著落后。初步判断为CPU侧调度问题，可能由上下文切换或CPU抢占引起。

2. WithoutStack场景：AllReduce通信成为性能瓶颈

关闭调用栈后，大部分算子执行趋于均衡，不过AllReduce算子存在一定的卡顿现象，1号卡落后0号卡约50ms。进一步分析两卡的Free时长发现：

• 0号卡Free时长：100ms
• 1号卡Free时长：50ms

两者差距约50ms，与AllReduce耗时高度吻合，表明快慢卡问题主要由该通信算子引发。结合调用栈分析，发现AllReduce前的两个index_select算子之间存在大段空白，疑似存在主机侧等待或内存迁移阻塞。

深入排查模型后，可能原因为显存数据向CPU迁移（mem搬迁） 导致的同步等待，进而引发后续算子执行延迟，形成“卡顿链”。

优化方案：算子替换与执行路径调优

1. 算子路径优化：从ACLOP到ACLNN

原始模型中Conv2D算子默认走ACLOP路径，性能会受一定影响。尝试通过关闭JIT编译，使用ACLNN路径：

torch.npu.config.allow_internal_format = False
torch.npu.set_compile_mode(jit_compile=False)

该改动使Conv2D算子性能提升超190倍。但由于该配置可能影响其他算子行为，存在兼容性风险，故不作为最终方案。

2. 算子替换：Conv2D → Matmul 融合实现

为规避路径切换风险，采用算子替换策略，将原始Conv2D操作替换为matmul融合实现，完整代码如下：

import torch
import torch.nn.functional as F
import numpy as np

# -------------------------
# 参数
# -------------------------
input_shape = [1, 3, 728, 728]
weight_shape = [1792, 3, 2, 2]
bias_shape = [1792]

stride = (2, 2)

# -------------------------
# 构造数据
# -------------------------
x = torch.randn(input_shape, dtype=torch.bfloat16)
w = torch.randn(weight_shape, dtype=torch.bfloat16)
b = torch.randn(bias_shape, dtype=torch.bfloat16)

N, Cin, H, W = x.shape
Cout, _, Kh, Kw = w.shape

Ho = H // Kh   # 52
Wo = W // Kw   # 52

# ==========================================================
# 1️⃣ 直接 reshape
# ==========================================================
# [N, Cin, H, W]
# → [N, Cin, Ho, Kh, Wo, Kw]
x_patch = x.view(
    N,
    Cin,
    Ho, Kh,
    Wo, Kw
)

# 调整维度顺序，把 patch 维度挪到一起
# [N, Ho, Wo, Cin, Kh, Kw]
x_patch = x_patch.permute(0, 2, 4, 1, 3, 5)

# 展平成 matmul 输入
# [N, Ho*Wo, Cin*Kh*Kw]
x_mat = x_patch.reshape(N, Ho * Wo, Cin * Kh * Kw)
# => [1, 2704, 588]

# ==========================================================
# 2️⃣ weight reshape
# ==========================================================
# [Cout, Cin, Kh, Kw] → [Cin*Kh*Kw, Cout]
w_mat = w.view(Cout, -1).transpose(0, 1)
# => [588, 1792]

# ==========================================================
# 3️⃣ BatchMatMul
# ==========================================================
bias = b.npu()

bias_expand = bias.view(1, 1, Cout).expand(N, Ho * Wo, Cout)
out = torch.baddbmm(
    bias_expand,
    x_mat.npu(),
    w_mat.npu().unsqueeze(0).expand(N, -1, -1),
    beta=1,
    alpha=1
)
# => [N, 2704, 1792]

# ==========================================================
# 5️⃣ reshape 成 conv 输出
# ==========================================================
out = out.transpose(1, 2).reshape(N, Cout, Ho, Wo)

该方案将原始Conv2D耗时从33ms降至可忽略级别，且快慢卡问题依然存在，说明性能瓶颈已从算子执行转移到数据迁移与通信调度层面。

优化结论与后续建议

1. 算子路径优化有效但存在风险：关闭JIT可显著提升Conv2D性能，但需评估对整体模型稳定性的影响，建议仅在特定路径中使用。
2. 算子融合是关键优化手段：将Conv2D替换为Matmul融合实现，不仅提升性能，还增强了算子可调度性，是更安全、可复用的优化方式。
3. 通信与内存调度是核心瓶颈：AllReduce与index_select之间的空白调用栈表明，主机侧数据迁移与同步存在一定性能瓶颈。后续可从以下方向优化：
   • 优化数据布局，减少显存与CPU间频繁搬运；
   • 采用异步数据预加载机制，避免同步等待；
   • 在通信前增加数据预处理流水，提升整体调度效率。

总结

本案例通过Profiling精准定位Prefill阶段的快慢卡问题，结合算子路径分析与融合优化，成功将关键算子性能提升超100倍。后续应聚焦于数据流调度与内存管理，结合异步通信与流水线设计，进一步释放Atlas 800I A2平台的算力潜力，实现Prefill阶段吞吐率的全面跃升。