之前帮一个团队优化LLaMA-7B推理，-batch=8、单卡的时候延迟是4.2ms。我让他们改成-batch=16试试，结果延迟变成8.8ms——慢了整整一倍。

"batch翻倍，延迟应该不到4.2ms才对啊？"他们很困惑。

我看了一下GE的调度日志，发现问题了：FlashAttention和后面的MLP分配到了同一个stream，计算和内存搬运是串行的。batch大了之后，内存带宽成了瓶颈，前一个算子没搬完，后一个算子只能等着。

"你们这是一个消防员接力的故事，"我说，“一个人灭火，另一个人要等灭火的人把水管收好才能开始浇水。正确的做法是两个人同时行動——消防员灭火的时候，第二个人已经在铺水管了。”

GE的Streams机制就是这样：把算子分配到不同的stream，让计算和内存搬运同时进行消防员灭火的时候，第二个人已经在铺水管了。

今天从FlashAttention的视角，看看GE是怎么通过多streams并行，让昇腾NPU的的计算资源和内存带宽都被充分利用的。

Streams是什么

Streams是昇腾NPU的一个核心抽象——可以理解为"并行执行走廊"。每个stream有自己独立的指令队列，不同stream的指令可以交错执行，只要没有数据依赖。

单stream（串行）：
Stream0: [加载Q] → [FlashAttention] → [写回O] → [加载K] → [MLP]
总时间 = t_加载 + t_FA + t_写回 + t_加载 + t_MLP

多streams（并行）：
Stream0: [加载Q] → [FlashAttention] → [写回O] → ...
Stream1: [加载K] → [MLP] → ...
总时间 = max(t_load+FA+写回, t_load+MLP) // 两个stream的并行部分叠加

Streams的核心价值：让内存搬运（HBM读写）和计算（TensorCore）并行。FlashAttention的计算本身很快（O(N²)的矩阵乘），但加载Q/K/V和写回输出很慢。多streams可以把"搬数据"和"算数据"重叠起来。

FlashAttention的Streams分配策略

FlashAttention在GE里的典型streams分配：

# GE给FlashAttention分配的streams
streams = {
 "Stream0": [
 # 输入加载和算子计算放在同一个stream
 Op("Load_QKV"), # 加载Q/K/V到L1
 Op("FlashAttention"), # 计算注意力
 Op("Store_O"), # 结果写回HBM
 ],
 "Stream1": [
 # 下一层（MLP）或下一个token的数据加载
 # 可以和Stream0并行执行
 Op("Load_Next"),
 Op("MLP"),
 ],
 "Stream2": [
 # 反向传播或其他异步操作
 Op("GradientSync"),
 ]
}

为什么这样分配？

FlashAttention的计算流程：

1. 加载阶段：Q/K/V从HBM搬到L1 Buffer
2. 计算阶段：TensorCore算Q@K^T、Softmax、Q@V
3. 写回阶段：输出O从L1写回HBM

其中，加载和写回占用大量HBM带宽，计算占用TensorCore。两者互不干扰，可以并行。

依赖管理：Streams不是万能的

Streams能并行是有条件的——没有数据依赖的算子才能放不同stream。

FlashAttention和下一层MLP的依赖：

Input ──[FlashAttention]──→ O ──[LayerNorm]──→ norm_out ──[MLP]──→ output

这里的O是FlashAttention的输出，同时是MLP的输入。所以：

• FlashAttention → LayerNorm：强依赖（同stream）
• LayerNorm → MLPN：强依赖（同stream）

但可以这样优化：

Stream0: [FlashAttention] ──→ [LayerNorm] ──→ [MLP]
Stream1: ──→ [预取下一次的QKV] ──→ （准备好下一次的输入）

下一次迭代的QKV加载和当前迭代的MLP计算没有数据依赖，可以并行。

GE的依赖分析器会自动算出这个"可以并行"的窗口：

// GE依赖分析示意
DependencyAnalyzer analyzer(graph);

for (auto& edge : graph.edges) {
 auto producer = edge.producer; // 生产者算子
 auto consumer = edge.consumer; // 消费者算子
 
 // 如果producer的输出是consumer的唯一输入，且没有副作用
 if (edge.is_single_consumer && !has_side_effects(producer)) {
 // 可以并行：producer和consumer可以用不同streams
 // 但要保证producer先执行完，consumer才能读到正确数据
 analyzer.mark_parallelizable(producer, consumer);
 }
}

FlashAttention多Streams实战配置

在GE里调整FlashAttention的streams分配：

# 方法1：默认配置（GE自动分配）
ge.set_option("enable_stream_parallel", True)

# 方法2：手动指定分配策略
ge.set_stream_assignment({
 "FlashAttention": 0, // Stream0
 "LayerNorm": 0, // 同Stream0，强依赖
 "MLP gate": 1, // Stream1，和FlashAttention并行
 "MLP up": 1, # Stream1
 "MLP down": 0, # 回Stream0，等待gate+up的结果
})

# 方法3：自定义分发策略
class MyAssignStrategy(StreamAssignStrategy):
 def assign(self, op_graph):
 # 计算密集型算子 -> Stream0
 # 内存密集型算子 -> Stream1
 for op in op_graph.ops:
 if op.is_compute_intensive():
 assign_to_stream(op, 0)
 elif op.is_memory_intensive():
 assign_to_stream(op, 1)

不同batch_seq下的Streams配置

batch和seq_len不同，Streams优化的侧重点也不一样：

配置	batch	seq_len	瓶颈	推荐streams
小batch短序列	1	512	计算	单stream够用
标准推理	8	1024	内存	2 streams
大batch长序列	32	2048+	带宽	3-4 streams
streaming生成	1	动态	IO	3 streams

具体配置：

# 小batch短序列：不用多streams
config_small = {
 "num_streams": 1,
 "stream_strategy": "default",
}

# 标准推理：2 streams足够
config_normal = {
 "num_streams": 2,
 "stream_strategy": "compute_overlap",
 "overlap_ratio": 0.7, # 计算和内存搬运重叠70%
}

# 大batch长序列：3-4 streams
config_large = {
 "num_streams": 4,
 "stream_strategy": "pipeline",
 "pipeline_depth": 3, # 3级流水线
}

实测收益

在昇腾910上跑LLaMA-7B，测试不同streams配置的性能：

配置	batch	seq_len	延迟	吞吐	加速比
单stream	8	1024	4.2ms	1,250	1×
2 streams	8	1024	2.8ms	1,780	1.42×
3 streams	8	1024	2.2ms	2,100	1.68×
4 streams	8	1024	2.0ms	2,200	1.76×
单stream	32	2048	18.5ms	820	—
4 streams	32	2048	9.2ms	1,450	1.77×

规律：

• streams数量存在拐点：超过4个streams后提升变缓（昇腾910只有4个硬件队列）
• batch越大，streams收益越高：大batch的内存搬运更多，并行收益更大

实战踩坑

坑一：streams分配过度

设置了8个streams，但昇腾910只支持4个，多余的stream被GE忽略。

解决：查询硬件支持的streams数量

# 查询支持的streams数量
max_streams = ge.get_device_capability("max_streams")
print(f"Device supports {max_streams} streams")

坑二：伪并行

明明设置了多streams，但实际执行时间没变化。

排查：启用GE的执行分析

# 启用执行分析
ge.enable_profiling()

# 运行推理
output = model.forward(x)

# 查看timeline
timeline = ge.get_execution_timeline()
timeline.print()

# 输出示例：
# Stream0: [==========] 10ms
# Stream1: [==== ] 5ms ← 并行了但很短
# Stream2: [ ] 0ms ← 没用到

如果某个stream几乎空闲，说明依赖分析太保守，或者算子之间的依赖太紧密。

坑三：数据竞争

多streams同时访问同一个tensor，导致数据错误。

解决：添加显式同步点

# 添加同步点
ge.add_sync_point(
 producer="FlashAttention",
 consumer="MLP",
 sync_type="data_ready" # 数据就绪后通知
)

坑四：batch=1没收益

单推理时，streams优化反而带来额外开销（stream切花、依赖检查）。

解决：对batch=1的场景关闭streams

# 条件性启用streams
if batch_size > 1:
 ge.set_option("enable_stream_parallel", True)
else:
 ge.set_option("enable_stream_parallel", False)

Streams vs HCCL多卡通信

Streams是多streams是单卡内的并行，HCCL是多卡间的通信。

单卡：GE Streams并行
 卡0 Stream0: 计算 + 加载
 卡0 Stream1: 预取
 
多卡：HCCL通信
 卡0 ───all-reduce──── 卡1
 │ │
 ────all-reduce──── ────┘

FlashAttention在单卡内用Streams优化，在多卡间用HCCL通信优化。

总结

GE的Streams机制核心价值：让内存搬运和计算并行，让计算单元不再因为等数据而空闲。

FlashAttention的典型Streams分配：

• Stream0：计算（FlashAttention + 后处理）
• Stream1：加载下一层/下一个token的数据
• Stream2：梯度同步（训练场景）

Streams不是越多越好——昇腾910只有4个硬件队列，超过4个没收益。batch越小、seq越短，单stream可能更快。

排查性能问题时，别忘了看GE的调度日志。搜索"Stream"关键字，看看算子是怎么分配的、有没有并行、一共有几个stream在跑。

很多推理性能问题不在算子实现，而在调度层面——Streams配置对了，延迟能降40%。

意外收获：Streams机制对训练也有效。训练时的反向传播可以用单独的stream，和前向传播并行——前向算当前batch的时候，反向已经在算上一个batch的梯度。这就是所谓的"流水线并行"。