news 2026/11/4 8:40:37

【学习】mKernel 分析

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
【学习】mKernel 分析

本文档基于 arXiv 2609.13585（mKernel，2026-09-11 提交，UC Berkeley）论文

1. 论文与背景速览

mKernel是面向多 GPU、多节点的融合 kernel 库：把"计算 + 机内 NVLink 通信 + 机间 RDMA"三者放进同一个 kernel，在tile 粒度上重叠执行。

  • 问题：MoE 训练中通信占 forward pass 43.6%、端到端训练 32%；设备间通信最高占执行时间 47%。算力增速快于网络带宽（GB300 NVL72 机架 720 PFLOP/s，每 GPU 出机架仅单 NIC 400–800 Gb/s）。
  • 现状：现有融合 kernel（Flux、Comet、TileLink 等）几乎都局限在单 NVLink 域；双流stream重叠只在 kernel/chunk 边界释放数据，粒度粗。
  • 答案：mKernel 在 2×8 H200 集群上，GEMM+AllReduce 最高1.72x、Ring Attention 最高1.88x（对比 cuBLAS/FlashAttention+NCCL 未融合基线）。
  • 反直觉发现：GPUDirect Async（IBGDA）相对 host-assisted GPU-initiated 通信几乎无吞吐增益。

来源：arXiv 2609.13585 / HTML 全文

2. 实现细节拆解

2.1 硬件层级（论文 Table 1）

层级链路带宽（每 GPU）SM 角色传输机制
片内HBM3e4.8 TB/scomputeloads/stores、TMA
节点内NVLink + NVSwitch450 GB/s本地通信TMA、NVSwitch
机间（IB）ConnectX-750 GB/ssend/receiveNIC RDMA
机间（EFA）AWS EFA SRD50 GB/ssend/receiveNIC RDMA

机间带宽只有 NVLink 的约1/9—— 这是"不能把两层同等对待"的物理前提。

2.2 运行时架构（论文 Figure 2）

每 GPU 一个持久 kernel，thread block 按角色划分：

Compute 块 → 写完成 tile + readiness flag → Intra-node 块 → NVSwitch 交换/广播（同节点 8 GPU） → Inter-node send 块 → host 内存 ring buffer 发布 48B 命令 → host proxy 线程 → 批量提交 NIC（libibverbs，最多 8 命令/批） → NIC RDMA → rail peer（同索引 GPU） → 接收端 arrival flag → receive 块消费数据 ↑ 可选 on-GPU controller：读进度计数器 → 发布目标通信 block 数

2.3 三层粒度解耦（tile / chunk / batch）

  • tile：计算与机内传输粒度（128 行、16 token 等，随 kernel 而定）；
  • network chunk：机间传输粒度，独立定长。GEMM+AllReduce 将 4 个本地归约 tile 组成256 KiBchunk；per-chunk 计数器让最后完成的 tile 即发布就绪，不等整个 GEMM 结束；
  • submission batch：proxy 每连接聚合至多8 条就绪命令一次提交，但保留每条命令独立的完成通知；用 bounded polling window 支持部分批量，避免稀疏到达被拖延。

接收端反向映射：Dispatch+GEMM 加载 token 前，检查与其字节范围相交的每个512 KiBchunk 是否到达。

2.4 SM 分区与运行时自适应

  • 静态扫描 2–64 通信 SM：最差 AllReduce 分区比最优慢约25x；最优值跨负载从 2 到 64 SM 变化。
  • 自适应 controller：每个 block 在任务边界检查共享"目标通信 block 数"，不一致则切换角色。按 cycle counter 估算成本，将目标设为两角色预测完成时间相等：
n_s* = B · (R_s·C_s) / (R_p·C_p + R_s·C_s)
  • 数据依赖决定初始分配：通信供计算时，等待输入的 compute 块可临时协助通信。
  • 效果：21 配置几何均值1.18xvs 静态最优固定分区（免逐形状调优）。

2.5 完成语义（跨平台的关键适配）

平台保序性通知路径
ConnectX-7（IB RC）每连接有序data write → 同一连接 flag write；接收端 GPU 直接轮询 flag，无需接收端 proxy
AWS EFA（SRD）独立写无序RDMA write with immediate（immediate 值标识 chunk）→ 接收端 proxy 在 write 接收完成后发布 arrival flag

原则：到达信号必须标识"完成的 payload"，而非"已提交的请求"。

2.6 五个 kernel（论文 Table 3）

Kernel并行依赖机内机间chunk
AllGather+GEMMTP通信→计算每 shard 机内 multicast每 shard 每目的节点发一次128 行
GEMM+ReduceScatterTP计算→通信TMA atomic add每节点部分和交换2–32 tile
GEMM+AllReduceTP计算→通信NVSwitch 归约+广播每 GPU 发 1/8 输出4 tile（256 KiB）
MoE Dispatch+GEMMEP通信→计算TMA pulltoken buffer 复制到 rail peer16 token / 512 KB
Ring AttentionSP步内独立TMA store KVKV slice 每节点发一次128 行 KV tile

关键调度选择：① 机间传输尽早发起（AllGather+GEMM 启动即提交输入 shard）；② NVSwitch 承担本地归约/复制；③ 计算 tile 按输入可用性排序（本地→同节点→远程）。

2.7 实验结果汇总

指标结果
GEMM+AllReduce最高 1.72x（CX7）/ 1.53x（EFA）
Ring Attention最高 1.88x（CX7）/ 1.78x（EFA）
AllGather+GEMM1.34x（CX7）/ 1.41x（EFA），全尺寸胜出
MoE Dispatch+GEMMEFA 上 3.1–4.5x；胜过 DeepEP+DeepGEMM
自适应 SM 分区21 配置几何均值 1.18x vs 静态最优
IBGDA vs host-assisted几乎无增益（fence+doorbell 开销所致）

3. 创新点提炼（4+1）

  1. SM 专业化（SM specialization）：同一持久 kernel 内 thread block 分四类角色（compute/intra-node/inter-node send/receive），角色实现与资源分配解耦——调 SM 预算即重平衡，不改计算块 warp 布局。
  2. 层级化数据移动（hierarchical data movement）：rail-optimized 拓扑下机间只交换 rail peer，NVSwitch 承担本地归约/广播，使走机间慢速网络的数据最小化；三层粒度（tile/chunk/batch）独立定长，兼顾"尽早传输"与"摊薄开销"。
  3. 可移植 host-assisted GPU-initiated 通信：GPU 产命令 + host proxy 提交，直接 libibverbs 实现（不依赖 NCCL/NVSHMEM），同一 kernel 可跑 InfiniBand 与 AWS EFA——把"传输层差异"封装在完成语义适配器中。
  4. 运行时动态 SM 分区：on-GPU controller 用实测进度 + 剩余工作实时调分配，免去逐 kernel/逐形状的静态调优。
  5. （发现性创新）IBGDA 增益甚微：GPU 直驱 NIC 需每 chunk 一次 system-scoped fence + doorbell，被 proxy 的"重叠提交 + 批量"抵消——挑战了"GPUDirect 一定更好"的默认假设。

4. API 与编程模型下的定义

mKernel 没有公开发布通用 API，但论文设计隐含了一套可形式化的多节点融合 kernel 编程模型。下面把它的原语与最小接口显式定义出来——这是把它移植到其他芯片（昇腾）的前提。

4.1 抽象出的编程原语

原语语义mKernel 实现
sm_role(block_id) ∈ {COMPUTE, INTRA, SEND, RECV}角色化 block 分配block index 静态分配 / 任务边界动态切换
ready_flag(tile_id)tile 就绪发布GPU 内存 flag + 内存序
chunk(ids, size)机间传输聚合单元4 tile→256 KiB 等
rail_peer(gpu_idx)机间交换对象相同本地索引的远程 GPU
command{peer, offset, len, chunk_id}传输命令48B、host 内存 ring buffer
completion_adaptor(transport)完成语义适配IB:>4.2 最小接口定义（示意伪代码，非论文原文）
// 多节点融合 kernel 抽象接口（基于 mKernel 设计的还原定义）template<typename ComputeKernel>voidfused_multi_node_kernel(ComputeKernel compute,// 计算 tile 逻辑RoleMap roles,// block → {COMPUTE|INTRA|SEND|RECV}intsm_budget[],// 各角色 SM 数（可被 controller 改写）TileToChunk group,// tile → chunk 聚合规则（独立定长）CompletionAdaptor notify,// 按传输平台选择的完成语义RailTopology rail)// rail peer 映射{// 1) 计算块产出 tile → 置 ready flag// 2) INTRA 块经本地高速网（NVSwitch/HCCS/UB）归约或广播// 3) SEND 块把就绪 chunk 发布为命令// 4) proxy 批量提交（可部分批量、背压）// 5) RECV 块按完成语义等待 payload 完成// 6) controller 读进度 → 写 sm_budget[] → 角色切换}

4.3 与现有编程模型的定位关系

模型层级与 mKernel 的关系
CUDA / Ascend C单核 kernel 编程mKernel 是"kernel 之上"的编排层
NCCL / HCCL集合通信库（kernel 外）mKernel 把集合下沉进 kernel（tile 粒度）
ThunderKittens / ParallelKittens机内融合原语mKernel 复用其 compute 原语，扩展机间
Triton-distributed / TileLink编译原语暴露mKernel 用显式 schedule，不依赖编译器
NVSHMEM / UEP / MSCCL++通信接口mKernel 借鉴 UEP 的 GPU-command/CPU-proxy 设计，但走裸 verbs

一句话定义：mKernel 的编程模型 =持久 kernel + 角色化 SM 预算 + 三层粒度解耦 + 传输无关的完成语义适配。

5. 昇腾芯片落地可行性分析

5.1 硬件层级映射（对照论文 Table 1）

mKernel 层级昇腾对应带宽量级对比
HBM3e 4.8 TB/s（片内）HBM 3.2 TB/s（Atlas 800I A3 整机）同级（片级需以 die 为单位核）
NVLink+NVSwitch 450 GB/s（节点内）HCCS 784 GB/s（8 卡整机高速互联）同级或更高
机间 IB 50 GB/s灵衢 UB 392–800 GB/s（超节点）/ RoCE（跨超节点）UB 显著更高
SM（compute/comm 角色）AI Core（AIC/AIV）+ AI CPU + DMA 引擎需重新设计角色分配
NIC RDMA verbsRoCE 驱动 / UB 内存语义（Load/Store）语义不同（见 5.3）
IBGDA（GPU 直驱 NIC）950 硬化集合通信加速单元昇腾已走"硬件专用通信单元"路线

5.2 逐创新点可行性评估

创新点昇腾可行性依据与改造点
1. SM 专业化（角色分工）高Ascend C 天然多核（block_idx 区分实例），可把部分 AI Core 分配为通信角色。但注意：AIC（矩阵）/AIV（向量）分工固定，通信角色更适合放在AI CPU / DMA 引擎或专用通信核；950 的硬化集合通信单元直接承担该角色——从"SM 分工"升级为"专用硬件分工"，效果更彻底。
2. 层级化数据移动（NVSwitch 归约→rail peer）高HCCS 是节点内一致性总线（天然支持归约语义），灵衢 UB 提供统一全局内存寻址 + 内存语义，比 IB 的消息语义更接近 mKernel 的"机内 peer memory"理想；HCCL 已有拓扑适配的 Pairwise 算法可复用。
3. host-assisted 命令队列 + 完成语义适配中高昇腾机间存在两条语义线：RoCE（消息语义，可复用 mKernel 的 proxy+完成适配器）与 UB（内存语义/Load-Store 同步通信，几乎不需要命令队列——同步语义天然保证>5.3 关键差异点：UB 的内存语义是"简化"还是"陷阱"？
  • 有利面：灵衢 UB 是总线级协议（Load/Store 同步通信、统一全局内存寻址），机间通信在语义上等同机内 peer memory——mKernel 花费大量设计去桥接的"内存语义 vs 消息语义"鸿沟（C2）在 UB 域内部分消解：tile 就绪即可被远程 Load，无需显式 RDMA 命令。
  • 需要警惕：① 保序性——mKernel 强调"arrival flag 不得先于 payload 可见"；UB 的内存语义若像 IB RC 一样保序则用 flag 轮询，若提供多路径/乱序（如超节点内多路径转发）则必须引入 mKernel 的 EFA 式完成适配（immediate + 代理 flag）；② 一致性——UB 跨 384 卡的全局内存寻址，其一致性模型（何时远程可见）决定 readiness flag 的设计；③ 流量控制——超节点内 800 GB/s 柜间带宽 + 多路径，背压语义需显式定义。
  • 结论：mKernel 的完成语义适配器应原样成为昇腾数据面的标准层，但适配对象从"IB vs EFA"变为"UB vs RoCE"。

5.4 分阶段落地方案建议

  1. 阶段一（库层，1–2 个季度）：在 HCCL 之上或内部实现 mKernel 式tile 级融合的通信原语（GEMM+AllReduce、Ring Attention 两个高收益 kernel 优先），复用 HCCL 拓扑与 Pairwise 调度；完成语义适配器先行（UB 同步 / RoCE proxy 两条路径）。
  2. 阶段二（kernel 层）：在 Ascend C 中定义"角色化 block + 运行时 SM 预算"扩展（若 CANN 支持动态 block 重映射），实现自适应分区；950 上优先对接硬化集合通信单元，把通信从 AI Core 完全卸载。
  3. 阶段三（硬件协同）：与 950 后续芯片协同设计"批量 doorbell / 网内归约（in-network reduction）"，把 mKernel 的层级化归约再下沉到 UB 交换机/UBFM 层面——与灵衢开放协议生态对齐。

6. 影响分析

6.1 对 kernel 编程模型的影响

  • 范式迁移：从"计算 kernel 与集合通信分置"转向"持久 kernel 内角色化编排"——通信不再是 kernel 外的一次调用，而是 kernel 内的一个角色。ParallelKittens（机内）→ mKernel（机间）→ 昇腾硬化通信单元（硬件化），是同一趋势的三级跃迁。
  • DSL 吸收：CuTeDSL、Ascend C 等将吸收"tile 就绪/消息就绪/角色预算"作为一等原语；编译器（Triton-distributed/TileLink 路线）与显式 schedule（mKernel 路线）将长期并存，最终收敛为可编译的多节点融合 kernel 语言。

6.2 对集合通信库的影响

  • NCCL/HCCL 的演进方向被验证：device-side API、0-SM 集合、网内归约（NCCL GIN/CFT、HCCL Pairwise）与 mKernel 的"kernel 内集合"是同一条路；未来集合通信库可能收缩为完成语义与拓扑的提供者，而融合逻辑上移到算子层。
  • 三层粒度解耦（tile/chunk/batch）可能成为集合通信的新成本模型维度——比单纯按消息大小建模更精细。

6.3 对硬件设计的影响

  • IBGDA 结论的行业警示：GPU/NPU 直驱网卡需要批量 doorbell 与硬件 fence 消除，否则收益被 per-operation 开销吞没；**硬化的集合通信加速单元（昇腾 950 已走）**与 NVIDIA 的网内归约（in-network reduction/multicast memory）是两条被验证的方向。
  • 完成语义（data-before-flag vs immediate 完成）应成为网络/总线协议的标准属性声明，而非各库各自适配。

6.4 对多路径 / UB 数据面的影响