CUDA_DEVICE_MAX_CONNECTIONS

发表于：
字数： 1039 阅读：≈ 3分钟

环境变量解释

CUDA_DEVICE_MAX_CONNECTIONS 这个环境变量在 Megatron-LM 源码中经常见到，主要是用来控制硬件调度 kernel 的顺序，以尽可能提高通信计算重叠场景的重叠率。

英伟达文档 cuda-environment-variables 说明了 CUDA_DEVICE_MAX_CONNECTIONS 默认为 8，可以设置为 1 到 32 的整数值，表示计算和拷贝队列的数量。（同样还有 CUDA_DEVICE_MAX_COPY_CONNECTIONS 表示对拷贝队列的设置（优先级高于 CUDA_DEVICE_MAX_CONNECTION）。）论坛回复 指出这个环境变量表示有多少个硬件队列和 CUDA 流发生映射关系。（如果这个值不够大则流之间存在虚假依赖。这样一来，设置为 1 就能完全保证 kernel 执行顺序和发起顺序一致。设置为 32 则尽可能设法让不同的流并发执行 kernels。）

Megatron-LM 1F1B + EP A2A overlap

Megatron-LM 2025.8.1 的提交支持了 1F1B + EP A2A overlap，也对 CUDA_DEVICE_MAX_CONNECTIONS 环境变量有描述。

代码：

思考：

  1. 警告信息中暗示 Blackwell 可能更好处理同时开启 TP/CP 和 EP overlap 的情况。这个或许和 Blackwell 的 cluster launch control 功能有关。
  2. 对通信计算重叠使用 CUDA_DEVICE_MAX_CONNECTIONS=32，是说明两个一前一后的 microbatch 中各通信、计算 kernels 并不是精心重叠的、需要依赖硬件的并发调度提高效率吗？

Tip

1F1B + EP A2A overlap 的发展：

  1. 2025.3 小红书和英伟达中国发 博客 介绍了这个方案。
  2. 2025.6 dots.llm1 使用了这一方案。
  3. 2025.7 Kimi K2 使用了这一方案。
  4. 2025.8 Megatron-LM 将这个方案 提交 到主分支。

也可以看出来 Megatron-LM 在流水线调度问题上始终不愿放弃 1F1B，就像 CP 方案不愿接入 Ulysses SP 一样。可能是代码上有较多对 1F1B 的依赖，比如 TP 反向的 bulk overlap 等。

知乎 - Ring Attention 性能问题引发的计算通信 overlap 分析

https://zhuanlan.zhihu.com/p/706805407

这也是一篇很有质量的分析通信计算重叠的文章。文章给出了一个默认状态下的通信计算例子（使用 CUDA_DEVICE_MAX_CONNECTIONS 默认值 8）：

可以看到只要通信 kernel 比计算 kernel 晚发起，两者就不能很好重叠，尽管它们之间没有依赖关系。

作者还尝试了不用 CUDA_DEVICE_MAX_CONNECTIONS 环境变量，而是手动在两个流之间添加同步，使得计算比通信晚发起，结果 overlap 效果也变好了，达到了相似的效果。

（图片来自文章）

为什么无手动同步时，也经常观察到通信 kernel 等计算 kernel 呢？我认为：

  1. 如果发起的 torch 集合通信操作使用到了某个输入张量，需要同步该张量所在的流。DeepWiki 链接 指出了代码中的显式同步操作。
  2. 计算流 kernel 先发起，占满了资源，等到计算 kernel 要结束的时候资源才腾出来。分析：光是 1 还不能解释上面 Perfetto 界面显示 GEMM kernel 快要结束、但还没结束时，通信 kernel 就已经开始执行的情况。我相信代码里面肯定是先发起了通信操作，再发起计算操作的，因此通信操作的事件同步点在计算操作之前，不会存在阻塞问题。但由于没有设置 CUDA_DEVICE_MAX_CONNECTIONS=1，先发起的通信 kernel 没有先调度，导致计算 kernel 占满了 SMs，只有快结束的时候通信 kernel 才得以运行。