CUDA stream创建和依赖，多流并行

cuda不同的stream是独立执行的，如何建立多个steam的分叉和合并依赖？

CUDA 不同 stream 默认是彼此独立调度的，要建立多个 stream 之间的"分叉 / 合并"依赖，最常用的机制就是：

cudaEventRecord() + cudaStreamWaitEvent()

可以把它理解成 CUDA stream 之间构建 DAG 依赖关系的基本原语。

假设想构造：

复制代码

                 ┌─ stream1: kernel_B ─┐
stream0: kernel_A                     ├─ stream3: kernel_D
                 └─ stream2: kernel_C ─┘

             fork                   join

也就是：

复制代码

A
├── B
└── C
    ↓
必须等 B、C 都完成
    ↓
D

实现方式如下：

复制代码

cudaStream_t s0, s1, s2, s3;
cudaStreamCreate(&s0);
cudaStreamCreate(&s1);
cudaStreamCreate(&s2);
cudaStreamCreate(&s3);

cudaEvent_t eventA;
cudaEvent_t eventB;
cudaEvent_t eventC;

cudaEventCreate(&eventA);
cudaEventCreate(&eventB);
cudaEventCreate(&eventC);


// 1. A 在 stream0 上运行
kernel_A<<<grid, block, 0, s0>>>();

// A 完成后记录 eventA
cudaEventRecord(eventA, s0);


// 2. fork：stream1 / stream2 都等待 A
cudaStreamWaitEvent(s1, eventA, 0);
cudaStreamWaitEvent(s2, eventA, 0);

kernel_B<<<grid, block, 0, s1>>>();
kernel_C<<<grid, block, 0, s2>>>();


// 3. B、C 分别记录完成事件
cudaEventRecord(eventB, s1);
cudaEventRecord(eventC, s2);


// 4. join：stream3 同时等待 B 和 C
cudaStreamWaitEvent(s3, eventB, 0);
cudaStreamWaitEvent(s3, eventC, 0);

kernel_D<<<grid, block, 0, s3>>>();


// 最后才需要 CPU 等待整个计算结束
cudaStreamSynchronize(s3);

核心就是下面两个模式。

1. Stream 分叉 Fork

例如：

复制代码

stream0
   A
   |
 eventA
  /   \
 ↓     ↓
s1     s2
 B      C

代码：

复制代码

kernel_A<<<..., s0>>>();
cudaEventRecord(eventA, s0);

cudaStreamWaitEvent(s1, eventA);
cudaStreamWaitEvent(s2, eventA);

kernel_B<<<..., s1>>>();
kernel_C<<<..., s2>>>();

注意 cudaStreamWaitEvent() 不会让 CPU 阻塞。它只是向对应 stream 插入一个 GPU-side dependency：

复制代码

s1:
wait(eventA) → B

s2:
wait(eventA) → C

因此 CPU 可以继续提交后续任务。

2. Stream 合并 Join

比如：

复制代码

s1: B → eventB ─┐
                 ├→ s3: D
s2: C → eventC ─┘

代码：

复制代码

kernel_B<<<..., s1>>>();
cudaEventRecord(eventB, s1);

kernel_C<<<..., s2>>>();
cudaEventRecord(eventC, s2);

cudaStreamWaitEvent(s3, eventB);
cudaStreamWaitEvent(s3, eventC);

kernel_D<<<..., s3>>>();

多个：

复制代码

cudaStreamWaitEvent(s3, eventB);
cudaStreamWaitEvent(s3, eventC);

实际上就是逻辑上的：

复制代码

D depends on B && C

只要 B 或 C 有一个没有完成，D 都不会执行。

特别重要的一点是：同一个 stream 内本身天然有顺序依赖，不需要 event。

例如：

复制代码

kernel_A<<<..., s0>>>();
kernel_B<<<..., s0>>>();
kernel_C<<<..., s0>>>();

天然保证：

复制代码

A → B → C

所以 event 主要用于建立：

复制代码

跨 stream dependency

因此整个 CUDA 执行可以理解成一个 DAG：

复制代码

        B
      ↗   ↘
A              E
      ↘   ↗
        C → D

例如可以写成：

复制代码

// A
A<<<..., s0>>>();
cudaEventRecord(eA, s0);

// A -> B
cudaStreamWaitEvent(s1, eA);
B<<<..., s1>>>();
cudaEventRecord(eB, s1);

// A -> C
cudaStreamWaitEvent(s2, eA);
C<<<..., s2>>>();

// C -> D
D<<<..., s2>>>();
cudaEventRecord(eD, s2);

// B,D -> E
cudaStreamWaitEvent(s3, eB);
cudaStreamWaitEvent(s3, eD);

E<<<..., s3>>>();

CUDA runtime 最终看到的就是：

复制代码

           ┌── B ─────┐
           │           ↓
A ─────────┤           E
           │           ↑
           └── C ── D ─┘

这其实也是理解 CUDA Graph 的一个很好方式。

还有一个实践上很关键的区别：

复制代码

cudaEventSynchronize(event);

和：

复制代码

cudaStreamWaitEvent(stream, event);

不是一回事。

前者：

复制代码

GPU
 ↓
CPU 等待
 ↓
CPU 再提交 GPU 工作

会让 host 参与同步，类似：

复制代码

GPU → CPU → GPU

而后者：

复制代码

GPU stream1
     ↓
   event
     ↓
GPU stream2

是纯 GPU-side dependency：

复制代码

GPU → GPU

所以在高性能 pipeline 中，通常优先使用：

复制代码

cudaEventRecord()
cudaStreamWaitEvent()

而尽量避免中间频繁：

复制代码

cudaDeviceSynchronize()
cudaStreamSynchronize()
cudaEventSynchronize()

因为这些会让 CPU 等待，从而降低异步流水能力。

你可以把几种同步关系简单记成：

依赖关系	方法
同一 stream 内 A → B	天然保证
stream A → stream B	`cudaEventRecord + cudaStreamWaitEvent`
1 个 stream 分叉到 N 个 stream	N 个 stream 都 wait 同一个 event
N 个 stream 合并到 1 个 stream	目标 stream wait N 个 event
CPU 等 GPU	`cudaStreamSynchronize / cudaEventSynchronize`
所有 GPU 工作全部完成	`cudaDeviceSynchronize`

所以，如果你在实现类似 LLM inference 的多 stream pipeline，比如：

复制代码

                 Attention
                /
Previous Layer
                \
                 MoE / Communication

                    ↓
                Next Layer

很自然就可以用 event 实现：

复制代码

             event_prev
              /      \
           stream1  stream2
             ↓        ↓
         compute     comm
             ↓        ↓
          event1    event2
              \      /
               stream3
                  ↓
              next layer

这种 event + stream wait 本质上就是 CUDA 中实现 fork-join DAG 的标准方法。

vLLM 和 SGLang 都已经把"多 CUDA Stream + Event 建立 fork/join 依赖"作为实际的性能优化手段使用。其中 vLLM 甚至已经把它抽象成了通用工具函数；SGLang 则在 MoE、KV Cache、通信/计算 overlap 等路径里直接使用。

vLLM：已经抽象成通用的 fork/join 工具

当前 vLLM 有一个专门的：

复制代码

vllm/utils/multi_stream_utils.py

https://github.com/vllm-project/vllm/blob/main/vllm/utils/multi_stream_utils.py

里面的 maybe_execute_in_parallel() 几乎就是我们上一条讨论的模式。它的注释明确说明：一个函数运行在当前 stream，另一个运行在 auxiliary stream，并通过 CUDA Event 同步。代码逻辑是：

复制代码

event0.record()

result0 = fn0()

with torch.cuda.stream(aux_stream):
    event0.wait()
    result1 = fn1()
    event1.record()

event1.wait()

也就是：

复制代码

                   ┌── current stream: fn0 ─────────┐
previous work ─────┤                                 ├── next work
                   └── aux stream:     fn1 ─────────┘
                       ↑                         │
                    event0                    event1

这正是：

复制代码

              fork
               │
         ┌─────┴─────┐
         ↓           ↓
       stream0     stream1
         fn0         fn1
         │           │
         └─────┬─────┘
               ↓
              join

而且 vLLM 还有更一般化的 execute_in_parallel()，支持 1 个主 stream + N 个 auxiliary streams 。源码注释直接写的是：start_event 从当前 stream fan out 到所有 auxiliary stream；每个 auxiliary stream 完成后记录自己的 done_event，最后主 stream 等待所有 done_event 完成，即 fan-in/join。

所以你之前问的：

复制代码

一个 stream → 多个 stream → 一个 stream

vLLM 当前实际上已经做成了一个通用 primitive。

vLLM 里的具体应用也很典型。例如当前 DeepSeek-V4 Attention 路径会把 Query/KV 相关计算留在默认 stream，同时把 indexer、MLA compressor 放到 auxiliary streams，并通过 execute_in_parallel() 并行执行。也就是说类似：

复制代码

                           ┌→ sparse indexer ──┐
LayerNorm / hidden state ──┼→ MLA compressor ─┼→ Attention
                           └→ Q/KV projection ─┘

官方代码明确注释：

Q projection 和 KV insertion 在 default stream；indexer 和 MLA compressor 使用 auxiliary streams。

Kimi K3 路径甚至更加直观。vLLM 会把：

复制代码

router gate

和：

复制代码

routed expert down projection

放在两个 stream 上执行：

复制代码

hidden_states
     │
     ├───────────────┐
     ↓               ↓
 Router Gate       Down Projection
 default stream     aux stream
     │               │
     └───────┬───────┘
             ↓
         MoE experts

源码注释明确说：

router gate 和 down projection 都只读取 hidden_states，因此 gate 在 default stream，down projection 在 aux stream，通过 maybe_execute_in_parallel join。

甚至还有 计算和通信 overlap。Kimi K3 的 LatentMoE 路径中，vLLM 会把：

复制代码

up-projection GEMM

和：

复制代码

tensor parallel all-reduce

通过 maybe_execute_in_parallel() 放到不同 stream：

复制代码

                  ┌── GEMM ──────────────┐
MoE output ───────┤                       ├── add
                  └── TP AllReduce ──────┘

即典型的：

compute stream + communication stream overlap

SGLang：MoE 中有非常标准的双 Stream fork/join

SGLang 当前的 Qwen2-MoE、GLM4-MoE 等实现中也能直接看到。

例如 GLM4-MoE：

复制代码

current_stream = torch.cuda.current_stream()

self.alt_stream.wait_stream(current_stream)

shared_output = self._forward_shared_experts(hidden_states)

with torch.cuda.stream(self.alt_stream):
    router_logits = self.gate(hidden_states)
    topk_output = self.topk(hidden_states, router_logits)
    final_hidden_states = self.experts(...)

current_stream.wait_stream(self.alt_stream)

final_hidden_states += shared_output

画出来就是：

复制代码

                    fork
                     │
        ┌────────────┴─────────────┐
        │                          │
        ↓                          ↓
 main stream                  alt stream
 Shared Expert               Router Gate
      │                           ↓
      │                          TopK
      │                           ↓
      │                     Routed Experts
      │                          │
      └─────────────┬────────────┘
                    ↓
                   join
                    ↓
          shared + routed

其中：

复制代码

self.alt_stream.wait_stream(current_stream)

相当于建立：

复制代码

current stream → alt stream

的 fork 起点依赖。

而：

复制代码

current_stream.wait_stream(self.alt_stream)

就是：

复制代码

alt stream → current stream

的 join。

SGLang Qwen2-MoE 当前也有类似路径。

还有一个很有意思的例子：SGLang KV Cache 写入也会使用双 stream。

当前代码中，在满足特定条件时：

复制代码

current_stream = device_module.current_stream()

alt_stream.wait_stream(current_stream)

k_cache[indices] = k

with device_module.stream(alt_stream):
    v_cache[indices] = v

current_stream.wait_stream(alt_stream)

也就是把：

复制代码

K cache write
V cache write

拆到两个 stream：

复制代码

                    ┌── write K ──┐
input K,V ── fork ──┤              ├─ join → Attention
                    └── write V ──┘

SGLang 更重要的用途：通信与计算 overlap

对于大规模 MoE 推理，这其实比单纯"两个 GEMM 并行"更重要。

SGLang 的 EP 文档当前明确介绍了：

复制代码

TBO = Two-Batch Overlap
SBO = Single-Batch Overlap

其目标都是把：

复制代码

Attention / MoE Compute

和：

复制代码

Dispatch / Combine / All-to-All Communication

重叠起来。

例如 TBO 可以理解为：

复制代码

Compute stream:

Batch 0: Attention ───────── MoE ──────────────
Batch 1:             Attention ───────── MoE ───

Comm stream:

Batch 0:        Dispatch ───────── Combine
Batch 1:                    Dispatch ───────── Combine

当前 SGLang DP Attention 的代码注释甚至明确写了：使用一个独立 comm stream，返回 CUDA Event，从而让另一个 ubatch 的 Attention + MoE compute 在 compute stream 上运行，与当前 ubatch 的 gather/combine 通信重叠。

这就是大模型推理中最典型的：

复制代码

                 ┌────── Compute Stream ──────┐
dependency ──────┤                             ├── dependency
                 └────── Comm Stream ─────────┘

两个引擎可以这样理解

场景	vLLM	SGLang
多 CUDA stream	✅	✅
CUDA Event 建立依赖	✅	✅
1→2 fork/join	✅	✅
1→N fork/join	✅ `execute_in_parallel()`	✅ 多种 overlap 路径
MoE 分支并行	✅	✅
GEMM/GEMM overlap	✅	✅
Compute/Communication overlap	✅	✅，尤其 EP/TBO/SBO
KV Cache / copy overlap	✅	✅
CPU/GPU async overlap	✅	✅

其中 vLLM 当前的 multi_stream_utils.py 特别值得看，因为它几乎就是一个可以直接拿来教学的 CUDA fork/join 实现。

有一点也很重要：多 Stream 并不是越多越好。对于 LLM Decode，小 batch 下很多 kernel 没有把 GPU SM / Tensor Core / memory bandwidth 吃满，因此类似

复制代码

router ─────┐
            ├ overlap
projection ─┘

往往有收益；但在大型 Prefill GEMM 已经把 GPU 算力吃满时，再启动一个 GEMM 通常没有真正的并行资源，反而可能引入 stream/event 调度开销。所以你会发现 vLLM Kimi K3 的代码甚至根据 num_tokens 做 threshold，只在 token 数较小时启用 auxiliary stream。

因此对于 LLM inference，我会把多 Stream 的高价值场景概括成：

复制代码

                       CUDA Multi-Stream
                              │
              ┌───────────────┼───────────────┐
              ↓               ↓               ↓
        Compute/Compute  Compute/Comm     Compute/Copy
          overlap          overlap           overlap
              │               │               │
       Router + GEMM    MoE + AllToAll    GPU → CPU
       Shared/Routed    GEMM + AllReduce   KV transfer

尤其是 MoE + EP + DeepEP 场景，多 stream + event 已经不只是一个"小优化"，而是实现计算/通信 overlap 的核心基础设施之一。