Optimizing ML Concurrent Computation and Communication with GPU DMA Engines
Optimizing ML Concurrent Computation and Communication with GPU DMA Engines #
§1 TL;DR #
GPU 上 GEMM 与 collective 并发（C3）仅实现理想加速比的 21%。调度优先级与 CU 分区提升至 42%；将通信 offload 至 DMA 引擎（ConCCL），消除 compute 与 L1/L2 interference，达 72%，最高 $1.67\times$。
§2 痛点 / 方法 / 结果 #
Q1 痛点 #
ML 分布式训练与推理（data parallelism、FSDP、nano-batching）在 GPU 上产生计算与通信的并发执行需求（C3: Concurrent Computation and Communication）。当 GEMM kernel 与 collective communication kernel 同时在 GPU CU 上运行时，两者在 compute units、L1/L2 cache 和 HBM 带宽上产生严重的 mutual interference。在 AMD MI300X 8-GPU 平台上实测 15 个 C3 场景（源自 LLaMA-70B/405B 训练），baseline 仅实现理想加速比的 21%（平均 $1.13\times$），远低于理论上限 $1.6\times$–$2.0\times$。
Problem Formulation #
Objective: 最大化 C3 加速比
$$S = T_{\text{serial}} / T_{\text{concurrent}}$$
其中 $T_{\text{serial}} = T_{\text{GEMM}} + T_{\text{comm}}$，理想加速比为 $S_{\text{ideal}} = (T_{\text{GEMM}} + T_{\text{comm}}) / \max(T_{\text{GEMM}}, T_{\text{comm}})$。
Inputs / Outputs: 一对无数据依赖的 (GEMM kernel, collective kernel) → 调度策略 + CU 分配方案 + 通信路径（CU vs DMA） → 最小化 $T_{\text{concurrent}}$。
Assumptions:
- GEMM 和 collective 之间无数据依赖
- GPU 资源可部分分区（MI300X 的 CU reservation feature）
- DMA 引擎可通过 HSA API 显式调度（
hsa_amd_memory_async_copy_on_engine） - 通信为 intra-node（8-GPU fully-connected topology）
Q2 方法 #
三层递进优化：
- Schedule Prioritization (SP): 先调度资源需求小的 communication kernel（workgroup 数作 proxy），避免 GEMM 抢占过多 CU。21% → 42% of ideal。
- Resource Partitioning (RP): 利用 MI300X CU reservation 显式分配 CU 给通信流。~41% of ideal，与 SP 相当；二者叠加无额外收益。
- ConCCL (DMA Offload): 将 all-gather / all-to-all offload 到 14 个 SDMA 引擎（位于 IOD 层、L2 之下）。释放全部 304 CU 给 GEMM，同时消除 L1/L2 cache interference。66% → 72% of ideal（后者对 memory-bound GEMM 额外减 CU 以改善 cache）。
- Interference bound: 给定双 kernel 的 CU 和 bandwidth 需求，并发执行时间上界 $T_{\text{concurrent}} \leq f(T_{\text{GEMM}}, T_{\text{comm}}, \alpha_{\text{CU}}, \alpha_{\text{BW}})$
- Resource allocation optimality: CU 分区的最优性证明（当前仅有 heuristic，无最优性保证）
- Bandwidth contention model: DMA 与 GEMM 共享 HBM 时的 throughput degradation 解析模型
- 硬件: AMD MI300X Infinity Platform, 8× MI300X, fully-connected（7 bi-directional links/GPU, 64GB/s/link uni），单 GPU 304 CU + 192GB HBM + 256MB Infinity Cache + 14 SDMA
- 软件: ROCm, rocBLAS (GEMM), RCCL (baseline collective), rocprof (profiling)
- 测量: 15 次执行（6 warm-up + 9 measured），多 GPU stream 并发 launch
- Memory-bound GEMM 的反直觉加速: 减少 CU → 线程数降低 → L2 cache 争用减少 → 整体加速。ConCCL_rp 利用此效应对 memory-bound GEMM 减 8 CU（MI300X 上的最优值），额外贡献 6 个百分点（66% → 72%）。
- CPU orchestration 的隐含代价: DMA 由 CPU 通过 HSA API 编排（CPU → DMA queue in system memory → DMA engine），launch/sync 开销在 <32MB 时不可忽略（$4\times$ slower than RCCL）。仅 bandwidth-bound collective（≥128MB）适合 DMA offload。
- 代码开放性: 未开源。最接近的参考实现：RCCL（AMD collective library），MSCCL++（Microsoft GPU communication stack，含 DMA proxy channel）。
- 社区复现: 未见独立复现报告。ConCCL 使用 AMD 内部 HSA API 和 MI300X CU reservation feature。
- 生态影响: RCCL/NCCL 可在 bandwidth-bound all-gather/all-to-all 中引入 DMA offload path；PyTorch Async Tensor Parallelism 已开始利用 DMA 做 fine-grained C3。
- 局限: DMA 无算术能力 → 不支持 all-reduce（data parallelism 最核心 collective）；仅验证 intra-node 8-GPU；ConCCL 用 simple direct algorithm 而非 production-grade ring/tree/hierarchical algorithm。
核心技术壁垒: GPU 已有的 DMA 引擎位于 IOD 层（L2 cache 之下），offload 通信到 DMA 同时消除两种 interference：(1) compute interference（DMA 不占 CU），(2) L1/L2 cache interference（DMA 数据路径绕过 per-XCD cache）。这一 insight 需要对 GPU 芯片拓扑（XCD vs IOD 的物理分离）和内存层次结构的精确理解，以及暴露 per-engine 调度的底层 API。
Q3 结果 #
| 配置 | 平均加速比 | % of ideal | 备注 |
|---|---|---|---|
| c3_base | $1.13\times$ | 21% | GEMM 先调度，CU 自然分配 |
| c3_sp | ~$1.27\times$ | 42% | 通信 kernel 先调度 |
| c3_rp | ~$1.26\times$ | 41% | CU 显式分区 |
| ConCCL | ~$1.40\times$ | 66% | 通信 offload 到 DMA |
| ConCCL_rp | ~$1.46\times$ | 72% | DMA + memory-bound GEMM 减 8 CU |
All-to-all 受益最大：c3_base $1.05\times$ → ConCCL $1.43\times$（all-to-all 需 64 CU，interference 更严重）。最高单点加速比 $1.67\times$。
§3 架构 / 方法图 #
Paper's Figure 1, verbatim (caption: "Baseline C3 (left) and C3 with ConCCL via DMA offloads (right)").
左侧为 baseline C3：GEMM（橙色）和 communication（绿色）共享 GPU CU 和 cache，产生 compute + memory interference。右侧为 ConCCL：通信 offload 到 DMA 引擎（IOD 层），GEMM 独占全部 304 CU，DMA 数据路径绕过 per-XCD L1/L2 cache。
Method Core: CU-based → DMA-based Communication #
通信 kernel 本质是数据搬运，不需要算术计算。GPU 内置的 DMA 引擎可以独立完成这一工作，将 CU 完全释放给 GEMM。
| Before (c3_base) | After (ConCCL) | |
|---|---|---|
| 通信执行位置 | GPU CU（32–64 CU） | SDMA engines（14 个，IOD 层） |
| GEMM 可用 CU | $304 - N_{\text{comm}}$ | 全部 304 CU |
| L1/L2 cache interference | 有（共享 per-XCD cache） | 无（DMA 绕过 L1/L2） |
| HBM bandwidth interference | 有 | 仍有（DMA 与 GEMM 共享 Infinity Cache / HBM） |
| CPU orchestration 开销 | 无（GPU kernel launch） | 有（HSA API，<32MB 时 $4\times$ slower） |
| 支持的 collective | 全部 | 仅 all-gather / all-to-all（DMA 无算术能力） |
Paper's Figure 4, verbatim (caption: "C3 taxonomy").
C3 taxonomy 三层分类：① 计算与通信的相对长度（G-long / C-long / GC-equal，15% 阈值）；④ 相对幅度；⑤ kernel 子类型（compute-bound vs memory-bound GEMM，latency-bound vs bandwidth-bound collective）。该分类锚定了全部 15 个实验场景的分组和分析。
§4 作者证明 #
无形式化作者证明 — 仅实证。
论文不包含正式定理、引理或收敛性证明。所有结论基于 MI300X 平台上 15 个 C3 场景的 microbenchmark 实验。
符号表 #
| 符号 | 含义 |
|---|---|
| C3 | Concurrent Computation and Communication |
| $S_{\text{ideal}}$ | $(T_{\text{GEMM}} + T_{\text{comm}}) / \max(T_{\text{GEMM}}, T_{\text{comm}})$ |
| CU | Compute Unit（MI300X 共 304 个，8 XCD × 38） |
| XCD | Accelerator Complex Die（含 CU 和 4MB L2） |
| IOD | I/O Die（含 SDMA、Infinity Cache、HBM 接口） |
| SDMA | System DMA engine（14 个/GPU） |
| G-long / C-long / GC-equal | C3 taxonomy 类型 |
| cb / mb | compute-bound / memory-bound GEMM |
最接近的形式化内容 #
§V-C 的 resource partitioning heuristic 用 roofline model + slowdown lookup table 预测最优 CU 分配：对每种 CU 分配 $k$，用 slowdown factor $\sigma_{\text{GEMM}}(k)$ 和 $\sigma_{\text{comm}}(k)$ 缩放 roofline 时间，取使 $\max(T_{\text{GEMM}} \cdot \sigma_{\text{GEMM}}, T_{\text{comm}} \cdot \sigma_{\text{comm}})$ 最小的 $k$。Roofline 假设 70% 效率。该启发式在 24/30 场景命中最优 CU 分配，剩余 6 个场景损失 ≤1.5%。
期望的形式化保证 #
§5 实验与数据 #
实验平台 #
C3 场景 #
7 个 GEMM kernel（5 compute-bound cb1–cb5 + 2 memory-bound mb1–mb2，源自 LLaMA-70B/405B 训练 with 8192 tokens/iteration）× 2 种 collective（all-gather / all-to-all），组成 15 个 C3 场景（8 G-long + 5 C-long + 2 GC-equal）。其中 7 个来自实际 LLaMA FSDP 训练，8 个为 synthetic 补充以覆盖 taxonomy。
Kernel 资源需求 #
Paper's Figure 5, verbatim (caption: "(a) GEMM kernel slowdown due to loss of compute units (CUs)... (b) All-gather, (c) All-to-all kernel slowdown with specific # CUs assigned").
三个关键观察：(a) Memory-bound GEMM（mb1）在失去 32–48 CU 时反而加速（更少线程 → 更低 cache 争用），compute-bound GEMM（cb1–cb5）减速 17–27%。(b) All-gather 仅需 ~32 CU 即达满速。(c) All-to-all 需 ~64 CU。通信 kernel 仅占总 CU 的 10–20%，但共享导致不成比例的 interference。
软件优化效果 #
Paper's Figure 8, verbatim (caption: "Speedups for C3 scenarios under study with and without schedule prioritization and resource partitioning").
SP 将达成率从 21% 提至 42%（all-to-all: 0–13% → 27–46%; all-gather: 24–46% → 38–67%）。RP 达 41%。二者叠加（c3_sp_rp）无额外收益 → 解决同一瓶颈（communication kernel 的 CU starvation）。SP 实现最简单（仅改 CPU 端 kernel launch 顺序），是 production-ready 的低成本优化。
ConCCL 主结果 #
Paper's Figure 10, verbatim (caption: "C3 speedup with ConCCL").
ConCCL 在所有 C3 类型和 collective 类型上均显著优于 CU-based 方案。All-to-all 改善最显著（c3_base $1.05\times$ → ConCCL $1.43\times$），因为 all-to-all 的 CU 需求（64 vs 32）和 memory traffic 更高，DMA offload 消除的 interference 量更大。ConCCL_rp 对 memory-bound GEMM 额外减 8 CU（利用 cache 效应），再提 6 个百分点至 72% of ideal。
ConCCL Isolated 性能 #
ConCCL（简单 direct algorithm）在 ≥128MB 时与 RCCL 持平，但在 <32MB 时慢至 $4\times$（CPU 端 launch/sync 开销未摊销）。所有 C3 场景通信量 ≥128MB，故 isolated 性能差异不影响 C3 评估。
Training Recipe & Scale #
不适用。本文非训练方法论文，为 GPU 级 C3 性能 characterization 与优化。
Dataset Analysis #
不适用。本文 workload 为 microbenchmark（isolated GEMM + collective 的并发组合），非 data-driven 研究。
§6 论证链 #
| Step | 论点 | 证据 | 依赖 |
|---|---|---|---|
| 1 | C3 理想加速比 $1.1\times$–$2.0\times$（avg $1.6\times$），baseline 仅实现 21%（$1.13\times$） | Fig.7 ideal 投影 + Fig.8 c3_base 实测 | 无 |
| 2 | Interference 根因：通信占 32–64 CU + 共享 L1/L2/HBM；compute-bound GEMM 对 CU 损失敏感（17–27% slowdown），memory-bound GEMM 反而可加速 | Fig.5 CU 损失 vs slowdown + Fig.6 bandwidth 利用率 | Step 1 |
| 3 | SP（先调度小 kernel）和 RP（CU reservation）各自提升至 ~42%，叠加无额外收益 → 二者解决同一瓶颈（CU starvation） | Fig.8 c3_sp ≈ c3_rp ≈ c3_sp_rp | Step 2 |
| 4 | DMA 引擎位于 IOD 层（L2 之下），offload 通信到 DMA 同时消除 compute interference（释放全部 CU）和 L1/L2 cache interference（绕过 per-XCD cache） | MI300X 架构 Fig.2 + DMA 路径 Fig.3 | Step 2 |
| 5 | ConCCL PoC 在 ≥128MB 时 isolated 性能与 RCCL 持平 → DMA bandwidth 利用率可达 CU-based collective 水平 | Fig.9 ConCCL vs RCCL | Step 4 |
| 6 | ConCCL C3 达 66% ideal，ConCCL_rp 达 72% ideal → DMA offload 是 C3 优化的有效路线，剩余 gap 来自 HBM bandwidth 共享 | Fig.10 全景对比 | Step 3, 5 |
§7 实现 cross-reference #
[实现未公开]
ConCCL 被明确标为 proof-of-concept，论文未开源代码。核心 API 调用为 HSA runtime 的 hsa_amd_memory_async_copy_on_engine（ROCm ROCR-Runtime hsa_ext_amd.h）。ConCCL 实现使用 direct algorithm：每个 GPU 将数据直接写入所有其他 GPU（single step），由 CPU 编排 DMA transfer 调度。
核心技术壁垒 #
DMA offload 的关键在于识别 SDMA 引擎在芯片拓扑中的物理位置（IOD 层，L2 之下）及其对 interference 消除的双重效应。复制此方案需要：(1) 对目标 GPU 的 chiplet 拓扑（XCD vs IOD 的分离）有精确了解，(2) 有暴露 per-engine 调度的底层 API（AMD HSA 有 _on_engine，NVIDIA CUDA 目前无等价接口），(3) 处理 DMA 绕过 L2 带来的 cache coherence 问题——DMA 无法观察到前序计算在 L2 中的 dirty data，需要显式 flush 或 write-through。