【学习】mKernel 分析
本文档基于 arXiv 2609.13585mKernel2026-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。现状现有融合 kernelFlux、Comet、TileLink 等几乎都局限在单 NVLink 域双流stream重叠只在 kernel/chunk 边界释放数据粒度粗。答案mKernel 在 2×8 H200 集群上GEMMAllReduce 最高1.72x、Ring Attention 最高1.88x对比 cuBLAS/FlashAttentionNCCL 未融合基线。反直觉发现GPUDirect AsyncIBGDA相对 host-assisted GPU-initiated 通信几乎无吞吐增益。来源arXiv 2609.13585 / HTML 全文2. 实现细节拆解2.1 硬件层级论文 Table 1层级链路带宽每 GPUSM 角色传输机制片内HBM3e4.8 TB/scomputeloads/stores、TMA节点内NVLink NVSwitch450 GB/s本地通信TMA、NVSwitch机间IBConnectX-750 GB/ssend/receiveNIC RDMA机间EFAAWS EFA SRD50 GB/ssend/receiveNIC RDMA机间带宽只有 NVLink 的约1/9—— 这是不能把两层同等对待的物理前提。2.2 运行时架构论文 Figure 2每 GPU 一个持久 kernelthread block 按角色划分Compute 块 → 写完成 tile readiness flag → Intra-node 块 → NVSwitch 交换/广播同节点 8 GPU → Inter-node send 块 → host 内存 ring buffer 发布 48B 命令 → host proxy 线程 → 批量提交 NIClibibverbs最多 8 命令/批 → NIC RDMA → rail peer同索引 GPU → 接收端 arrival flag → receive 块消费数据 ↑ 可选 on-GPU controller读进度计数器 → 发布目标通信 block 数2.3 三层粒度解耦tile / chunk / batchtile计算与机内传输粒度128 行、16 token 等随 kernel 而定network chunk机间传输粒度独立定长。GEMMAllReduce 将 4 个本地归约 tile 组成256 KiBchunkper-chunk 计数器让最后完成的 tile 即发布就绪不等整个 GEMM 结束submission batchproxy 每连接聚合至多8 条就绪命令一次提交但保留每条命令独立的完成通知用 bounded polling window 支持部分批量避免稀疏到达被拖延。接收端反向映射DispatchGEMM 加载 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-7IB RC每连接有序data write → 同一连接 flag write接收端 GPU 直接轮询 flag无需接收端 proxyAWS EFASRD独立写无序RDMA write with immediateimmediate 值标识 chunk→ 接收端 proxy 在 write 接收完成后发布 arrival flag原则到达信号必须标识完成的 payload而非已提交的请求。2.6 五个 kernel论文 Table 3Kernel并行依赖机内机间chunkAllGatherGEMMTP通信→计算每 shard 机内 multicast每 shard 每目的节点发一次128 行GEMMReduceScatterTP计算→通信TMA atomic add每节点部分和交换2–32 tileGEMMAllReduceTP计算→通信NVSwitch 归约广播每 GPU 发 1/8 输出4 tile256 KiBMoE DispatchGEMMEP通信→计算TMA pulltoken buffer 复制到 rail peer16 token / 512 KBRing AttentionSP步内独立TMA store KVKV slice 每节点发一次128 行 KV tile关键调度选择① 机间传输尽早发起AllGatherGEMM 启动即提交输入 shard② NVSwitch 承担本地归约/复制③ 计算 tile 按输入可用性排序本地→同节点→远程。2.7 实验结果汇总指标结果GEMMAllReduce最高 1.72xCX7/ 1.53xEFARing Attention最高 1.88xCX7/ 1.78xEFAAllGatherGEMM1.34xCX7/ 1.41xEFA全尺寸胜出MoE DispatchGEMMEFA 上 3.1–4.5x胜过 DeepEPDeepGEMM自适应 SM 分区21 配置几何均值 1.18x vs 静态最优IBGDA vs host-assisted几乎无增益fencedoorbell 开销所致3. 创新点提炼41SM 专业化SM specialization同一持久 kernel 内 thread block 分四类角色compute/intra-node/inter-node send/receive角色实现与资源分配解耦——调 SM 预算即重平衡不改计算块 warp 布局。层级化数据移动hierarchical data movementrail-optimized 拓扑下机间只交换 rail peerNVSwitch 承担本地归约/广播使走机间慢速网络的数据最小化三层粒度tile/chunk/batch独立定长兼顾尽早传输与摊薄开销。可移植 host-assisted GPU-initiated 通信GPU 产命令 host proxy 提交直接 libibverbs 实现不依赖 NCCL/NVSHMEM同一 kernel 可跑 InfiniBand 与 AWS EFA——把传输层差异封装在完成语义适配器中。运行时动态 SM 分区on-GPU controller 用实测进度 剩余工作实时调分配免去逐 kernel/逐形状的静态调优。发现性创新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)机间交换对象相同本地索引的远程 GPUcommand{peer, offset, len, chunk_id}传输命令48B、host 内存 ring buffercompletion_adaptor(transport)完成语义适配IB:>4.2 最小接口定义示意伪代码非论文原文// 多节点融合 kernel 抽象接口基于 mKernel 设计的还原定义templatetypename ComputeKernelvoidfused_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 把集合下沉进 kerneltile 粒度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 1mKernel 层级昇腾对应带宽量级对比HBM3e 4.8 TB/s片内HBM 3.2 TB/sAtlas 800I A3 整机同级片级需以 die 为单位核NVLinkNVSwitch 450 GB/s节点内HCCS 784 GB/s8 卡整机高速互联同级或更高机间 IB 50 GB/s灵衢 UB 392–800 GB/s超节点/ RoCE跨超节点UB 显著更高SMcompute/comm 角色AI CoreAIC/AIV AI CPU DMA 引擎需重新设计角色分配NIC RDMA verbsRoCE 驱动 / UB 内存语义Load/Store语义不同见 5.3IBGDAGPU 直驱 NIC950 硬化集合通信加速单元昇腾已走硬件专用通信单元路线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–2 个季度在 HCCL 之上或内部实现 mKernel 式tile 级融合的通信原语GEMMAllReduce、Ring Attention 两个高收益 kernel 优先复用 HCCL 拓扑与 Pairwise 调度完成语义适配器先行UB 同步 / RoCE proxy 两条路径。阶段二kernel 层在 Ascend C 中定义角色化 block 运行时 SM 预算扩展若 CANN 支持动态 block 重映射实现自适应分区950 上优先对接硬化集合通信单元把通信从 AI Core 完全卸载。阶段三硬件协同与 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 路线与显式 schedulemKernel 路线将长期并存最终收敛为可编译的多节点融合 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 数据面的影响mKernel 的完成语义适配器 乱序/多路径传输的通用完成模型保序网用>6.5 对生态与标准的影响在灵衢 UB 开放协议2025-09 开放 2.0 规范、OCP 多路径规范等背景下kernel 内角色化通信有望成为跨厂商中间层上层 DSLAscend C/CuTeDSL与下层传输UB/IB/EFA/MRC之间需要一张 mKernel 式完成语义 层级化拓扑的标准接口——这是比单点性能提升更长期的生态价值。7. 结论mKernel 的价值把融合 kernel 从单 NVLink 域推进到多节点四原则SM 专业化/层级化数据移动/可移植 host-assisted 通信/运行时动态分区 一个反直觉发现IBGDA 增益甚微构成kernel 层通信编排的完整设计模板。昇腾落地判断核心思想角色分工、层级化归约、完成语义适配整体可行且天然契合——灵衢 UB 的带宽与内存语义甚至优于 mKernel 的测试平台主要工程缺口在 Ascend C 编程模型运行时角色切换、持久 kernel与 UB 一致性/保序模型的确认。最大协同点昇腾 950 的硬化集合通信加速单元 mKernel SM 分工做通信的硬件化终态完成语义适配器 三层粒度分层可直接沉淀为灵衢 UB 数据面的统一标准层。附主要外部来源mKernel 论文arXiv 2609.13585abs / html昇腾硬件与 CANN昇腾社区产品页Atlas 800I A3、昇腾社区 Asc-Comm、华为官方新闻灵衢发布、Atlas 900/950 超节点、arXiv 2506.12708CloudMatrix384、arXiv 2508.02520910C 架构、arXiv 2607.20120昇腾科学计算、CANN NEXT 950 架构详解鲲鹏昇腾社区、昇腾社区 HCCL Alltoall Pairwise 博客