DITRON是由字节跳动豆包大模型团队联合多所高校提出的一种分布式多级分块编译器旨在解决大型语言模型LLM分布式训练与推理中现有方案在灵活性、可扩展性和可移植性方面的不足。其核心贡献是将传统单设备分块编程模型如Triton扩展为支持分布式场景的三级抽象并通过编译优化实现计算与通信的高效重叠最终在多种硬件平台上取得超越专家调优库的性能。1. 研究背景与问题现状问题LLM规模急剧增长分布式执行成为刚需通信开销占总运行时间的20%~80%成为主要瓶颈高性能库CuBLAS/NCCL性能优异不可编程难以适配快速演进的模型架构现有张量编译器Triton等灵活性强缺乏对分布式内存层次的有效抽象可扩展性受限部分分布式编译器TileLink/Pallas初步探索通常仅支持单一粒度或节点内缺乏统一的多级抽象各硬件后端差异巨大现有框架移植成本高难以跨平台部署核心目标设计一种既灵活可编程、又可扩展支持多级并行、且可移植跨硬件的分布式编译器。2. 三大设计原则灵活的编程接口基于成熟的Triton分块模型做扩展而非从零设计新语言。现有单设备内核只需少量修改即可转为分布式版本。可扩展的编程范式采用多级分块思想——细粒度分块映射到高速互连如NVLink粗粒度分块映射到低速互连如PCIe/以太网以适应任意规模和异构集群。可移植的统一原语定义一套硬件无关的通信/同步原语新增硬件后端时只需实现对应的翻译规则即可。3. 系统架构三层设计DITRON采用前端-中端-后端三模块架构如下图所示3.1 前端三级编程接口核心创新层级粒度操作对象核心特性典型场景核心级Core-Level细粒度分块如128×128单GPU内计算静态形状继承Triton语法调用Tensor Core/TMAGEMM、Attention等计算逻辑设备级Device-Level粗粒度块Chunk跨GPU数据搬移动态形状支持MoE运行时路由DMA/NIC语义putmem/getmemAllToAll、跨节点数据传输任务级Task-Level整个模型/子图分布式集群调度将计算通信融合为MegaKernel消除内核启动开销保持硬件持续忙碌整网推理/训练这种三级抽象的核心价值在于将程序的逻辑视图用户写的分块计算与物理执行在不同硬件层次上的调度与通信彻底解耦。3.2 中端计算-通信交织Swizzle中端的关键优化是分布式交织Distributed Swizzling其核心思想是重排分块的执行顺序以最大化计算与通信的重叠Gather模式如AllGatherGEMM优先从远程获取数据提前发起fetch请求将本地HBM当作远程内存的缓存。Scatter模式如GEMMReduceScatter优先计算需要发送到最远节点的结果分块使长距离通信能尽早启动。具体实现中通过一个swizzle_func将原始分块ID映射为新的分块ID实现rank感知的偏移调度。论文还专门处理了非完美分块non-perfect tiling这一实际场景中的复杂情况。3.3 后端统一原语与代码生成定义硬件无关原语分为分布式原语wait/notify/rank等、SIMT原语sync/atomic等和SHMEM设备原语putmem/getmem等三类。通过LLVM CallExtern链接到厂商特定通信库如NVSHMEM/ROCSHMEM已为NVIDIA和AMD两个后端完整实现。支持分布式IR到LLVM IR的降级最终生成高效机器码。4. 支持的并行策略DITRON原生覆盖LLM训练/推理所需的所有主流并行范式并行策略对应的重叠内核张量并行TPAllGatherGEMM, GEMMReduceScatter, GEMMAllReduce序列并行SPGEMMAllToAll专家并行EPAllGatherMoE, MoEAllReduceDispatch/Combine流水线并行PP高效PP通信内核可灵活与其他层重叠5. 主要实验结果5.1 推理性能测试场景对比基线加速效果AG-GEMM8×H800CuBLASNCCL几何加速1.43×GEMM-RS8×H800CuBLASNCCL几何加速1.27×GEMM-AR8×H800CuBLASNCCL几何加速1.32×AG-MoE8×H800CuBLASNCCL几何加速19.18×MoE-AR8×H800CuBLASNCCL平均加速13.89×Attention模块Qwen3-32BCuBLASNCCLPrefill:1.12×Decode:1.26×FFN模块128k tokensCuBLASNCCLAllReduce:1.17×AGRS:1.27×端到端推理vLLM集成batch≥128vLLM5%~30%吞吐提升分布式MegaKernel单batchTorch Eager6.28×5.2 训练性能测试场景关键结果TP强扩展8~32 GPUs加速比 0.80×~1.71×GEMM变小是主要限制因素SP弱扩展8~128 GPUsGEMMAllToAll 始终保持一致性加速EP扩展512专家topk8Dispatch/Combine 加速1.04×~4.70×企业级训练部署MFU提升 10%月节省约50万GPU小时训练成本企业级推理部署端到端提升 20%边缘推理 30%5.3 跨平台可移植性平台对比基线加速效果AMD GPU (8×)RocmBLASRCCLAG-GEMM:1.11×GEMM-RS:1.16×模块级1.03×~1.07×PCIe GPU (8×)CuBLASNCCL几何平均8.33×多组不同工作负载6. 实际部署成果训练场景已在企业级训练中部署覆盖数十亿到数千亿参数模型AttentionSP加速20%MoE端到端提升10%优化器步骤加速20%。推理场景已部署于云服务推理TP和边缘推理机器人设备性能提升约20%~30%。开发效率相比专家手写CUDA重叠库如FLUX实现代码量减少一个数量级以上开发周期从数月缩短至数天且逐比特精度与原生实现一致。7. 总结DITRON的核心贡献可归纳为首次将分块级编译的思想Triton范式系统性地拓展到分布式多级硬件层次上通过三级编程抽象、计算-通信交织优化和统一后端原语实现了“像写单卡Triton一样写分布式内核”的目标性能不输甚至超越专家手写代码且天然支持多硬件平台。它本质上是一个编译层面的分布式编程民主化工具——降低了分布式内核开发的门槛同时通过编译优化保证了高性能填补了“高灵活性编译器”与“高性能专家库”之间长期存在的鸿沟。这里是自己的论文阅读记录感兴趣的话可以参考一下如果需要阅读原文的话可以看这里如下所示项目地址在这里如下所示摘要大型语言模型LLM的扩展目前受到分布式编程刚性不足的制约。虽然高性能库如 CuBLAS 和 NCCL提供了优化的原语但它们缺乏快速发展的模型架构所需的灵活性。相反现有的张量编译器无法有效应对分布式集群的复杂内存层次结构。为弥合这一差距我们提出了 DITRON一个可扩展的分块级编译器旨在使高性能分布式内核的开发大众化。DITRON 引入了一种新颖的分层编程抽象涵盖核心Core、设备Device和任务Task三个层级以将张量程序高效地映射到异构分布式硬件上。这种抽象使 DITRON 能够支持多种并行策略同时抽象化了节点间和节点内通信的复杂性。在大规模集群上的评估表明DITRON 的性能可与专家调优的 CUDA 库相当甚至更优在独立内核上实现了 6%−30% 的加速在 vLLM 的端到端推理中实现了 5%−30% 的加速。此外DITRON 展现了强大的可移植性在 NVIDIA 和 AMD 平台上均实现了显著加速。DITRON 已在企业级训练和推理任务中部署。它在训练任务中实现了超过 10% 的模型浮点利用率MFU提升每月可节省约 50 万 GPU 小时的训练成本。对于推理任务它提供了超过 20% 的端到端性能提升并已应用于云服务推理和边缘推理场景。1 引言大型语言模型LLM的快速发展给底层分布式系统带来了巨大压力。只有具备高速、可扩展的分布式执行能力才能将具有多样化新兴架构的模型实际部署 [2, 5, 7, 8, 21-23]。历史上优化工作呈现两极化像 Triton [27] 和 TileLang [26] 这样的编译器专注于单设备内核优化而分布式通信则依赖于像 NCCL [20] 这样的刚性库或特定领域的通信编译器 [24]。然而随着集群规模的扩大通信开销已成为主要瓶颈——在训练和推理中占运行时开销的 20% 到 80% [4]。这给研究人员带来了沉重负担他们现在必须具备双重能力既能设计新颖的算法又能将其实现为高度优化的分布式内核以证明其可扩展性。为应对这一挑战社区主要依赖两种方法分布式领域特定语言DSL编译器和分布式 CUDA 库。DSL 编译器 [10, 11, 33] 为描述并行程序提供了高级抽象而分布式库 [4, 20, 31, 32] 则提供了高度优化、专家调优的算子。然而这两种方法都存在明显的局限性。首先灵活性受到影响。分布式库本质上是不可编程的阻碍了对新架构的探索。同时DSL 编译器通常将用户限制在算子级粒度 [11, 28]缺乏进行细粒度优化的表达能力。其次可扩展性常常受限于刚性分块假设。像 Pallas [10] 和 TileLink [33] 这样的框架支持分块级编程但通常缺乏在大规模集群复杂层次结构中流水线执行所必需的统一多级分块抽象。第三缺乏可移植的原语使得将这些框架移植到新兴硬件后端成本高得令人望而却步。一个理想的编译器应在其设计中优先考虑可用性而不牺牲性能。但这一标准尚未被任何现有工作达到。具体而言我们认为分布式张量编程应遵循三个核心设计原则。首先它应提供灵活的编程接口与主流张量编译器的规范保持一致。正如近期报告 [14] 所述像 Triton [27] 这样的分块级编译器已超越 CUDA成为业界主流的核编程框架。扩展现有的分块编译器以支持分布式场景天然比从头开发新的编程语言更具优势。这种扩展应确保遗留内核能够以最少的代码修改转换为分布式版本。为此我们采用了 Triton 的分块编程模型并在此基础上实现了分布式扩展。其次它应体现一种可扩展的编程范式能够适应任意规模和问题形状的集群。不同任务的硬件配置差异巨大。一个分布式系统可能包含高带宽互连如 NVLink、硬件加速的交换机内计算能力如 NVSwitch以及低带宽链路包括 PCIe 和以太网。任何先验假设都可能使编译器在现实世界的系统部署中不切实际。为无缝适应此类异构系统我们主张在编程模型中支持多级分块其中细粒度分块映射到高速连接粗粒度分块映射到低速连接。第三它应提供一套统一的、可跨不同硬件后端移植的原语。不同的硬件后端具有不同的硬件拓扑和底层技术。因此低级编程模型需要截然不同的优化策略和编程范式。为实现对多后端无缝支持我们引入了一套硬件无关的原语。集成新的硬件后端可以通过实例化这些原语并实现相应的翻译规则来完成。在这些设计洞察的指导下我们提出了 Ditron一个分层分布式张量编译器。Ditron 结构分为三层前端、中端和后端。在前端Ditron 实现了三个层次的分块核心级分块、设备级分块和任务级分块。核心级分块使用户能够利用小型静态形状的向量或矩阵来调用硬件加速单元如张量内存加速器TMA和张量核心同时保持与现有 Triton 编程语言在功能和性能上的完全兼容。设备级分块释放了硬件 DMA/RDMA 引擎的潜力这些引擎针对高吞吐量的大块数据传输进行了优化。此分块级别采用动态形状并支持运行时形状计算这对于像混合专家模型MoE这样的动态模型架构至关重要。任务级分块更进一步实现了模型级分块它将整个模型融合到单个内核中然后将其部署并在分布式集群上执行以最大化硬件资源利用率。在中端Ditron 将算子级集体通信转换为分块级语义并实现了计算与通信的重叠。DITRON 原生支持 LLM 训练和推理的所有主要并行范式包括张量并行、序列并行、专家并行和流水线并行。集体通信原语如 AllGather、ReduceScatter、AllToAll 和 AllReduce被动态划分为块每个块映射到相应计算内核的一组分块。同时分布式同步事件被嵌入到计算内核中以实现与通信内核的无缝协调。图1 GPU 分布式内存层次结构与集群概览。在后端DITRON 提供了一套硬件无关的原语涵盖地址映射、数据访问和同步机制。由于这些原语与硬件无关只需为任何目标后端实现相应的翻译规则即可将它们转换为硬件特定的汇编代码。我们已经为 NVIDIA 和 AMD 后端实例化了这些翻译规则从而使 DITRON 能够支持这两个平台上的广泛 GPU。通过广泛的验证和评估DITRON 在各种工作负载下相对于供应商提供的非重叠内核能够实现平均 1.27× 至 19.18× 的加速甚至比用 CUDA 实现的专家调优重叠库快 6%−30%。当与 vLLM [13] 集成用于大批量推理时它也带来了 5%−30% 的端到端性能提升。在 AMD GPU 上相对于 RocmBLASRCCL 的加速比为 1%−38%在 PCIe GPU 上相对于 CuBLASNCCL 的平均加速比为 8.33×。2 背景与相关工作为理解 DITRON 背后的设计原理我们首先刻画现代分布式集群的硬件层次结构然后分析现有编程模型为何无法与该层次结构对齐。2.1 分布式硬件的层次结构分布式 GPU 集群的内存和计算层次结构本质上是非均匀的。如图 1 所示有效的分布式编程需要管理跨越三个不同领域的数据移动每个领域具有截然不同的带宽和延迟。在设备内部领域核心级涉及 HBM、L2 缓存和寄存器之间的数据移动。此处的优化依赖于最大化 SRAM 中的数据重用并利用张量核心等专用计算单元。此领域的程序通常由网格级、线程块级和线程束级的计算与通信指令组成。在横向扩展领域节点内现代 GPU [19] 节点通过 NVLink 等高带宽结构连接。此领域通过加载/存储指令支持统一内存访问UMA但也引入了 NUMA 特性远端访问延迟不可忽略。至关重要的是此领域提供了专门的硬件功能如用于网内归约的 NVLink Sharp [17]而标准软件常常忽视这些功能。在纵向扩展领域节点间节点间的通信依赖于以太网或 InfiniBand。此领域严格遵循 NUMA 架构数据传输需要通过直接内存访问DMA引擎和网络接口卡NIC进行显式协调。这里的带宽差距显著HBM 提供数 TB/s 的带宽而节点间链路通常运行在数十 GB/s。图2 Ditron 由前端接口、中端交织Swizzle以及后端原语和代码生成组成。分布式 LLM 工作负载的根本挑战在于这些层级并非孤立存在单个操作例如分布式矩阵乘法通常同时跨越所有三个领域。2.2 现有编程模型的局限性尽管硬件具有层次性现有的软件栈大多未能提供一个统一的抽象来捕捉这些细微差别。手工调优库先前的工作 [4, 9, 31, 32] 尝试通过提供用于计算-通信重叠的手工调优内核来弥补这一差距。例如像 FLUX [4] 和 COMET [31] 这样的库将通信原语直接集成到 CUTLASS 内核中。然而它们的编程接口通常晦涩难用需要深入了解汇编PTX或复杂的 C 模板元编程。这种刚性使得研究人员难以将这些内核适应新的模型架构。编译器和 DSL诸如 CoCoNet [11] 和 DistEinsum [28] 之类的编译器为分布式编程提供了算子级 DSL。最近像 Pallas [10]、TileLink [33] 和 IRIS [3] 这样的框架尝试将分块概念扩展到分布式环境。虽然很有前景但它们常常难以处理多级层次结构的复杂性。例如TileLink 主要关注横向扩展抽象但未能暴露高效纵向扩展通信所需的控制。3 Ditron 系统设计基于第 1 节讨论的洞察我们提出了 Ditron一个统一的编译器栈旨在弥合高级分布式算法与低级硬件异构性之间的鸿沟。如图 2 所示Ditron 采用模块化设计包括一个分层前端接口、一个以优化为中心的中端交织器swizzle以及一个带有统一原语的可移植后端代码生成器。3.1 前端三级编程接口Ditron 的核心创新在于将分布式张量程序的逻辑视图与其物理执行解耦。我们引入了一个三级分块抽象将不同的程序语义映射到相应的硬件领域。核心级接口在最细粒度上Ditron 继承了 Triton 的分块级语义以管理单个 GPU 内的计算。此级别操作于静态形状的分块例如128×128 的块以最大化张量核心和 TMA 等专用计算单元的利用率。通过保持与标准 Triton 语法的兼容性Ditron 允许用户无缝复用现有的优化内核来计算逻辑。我们在附录 D 中展示了实现高性能 AllGatherGEMM、GEMMReduceScatter 和 GEMMAllReduce 重叠内核的代码其中计算内核只需修改几行代码高亮显示即可用于分布式系统。核心级接口主要包含信号控制wait, notify、交换机上计算multimem以及分块交织tile swizzling如图 2 中的表格所示。设备级接口与单设备编译器不同Ditron 引入了设备级抽象来管理跨分布式内存层次结构的数据移动。此级别操作于块chunk上块是由多个细粒度分块组成的粗粒度数据块。此抽象设计有两个关键特性。首先我们使用以 DMA 为中心的语义。与核心级的加载/存储指令不同设备级原语例如putmem, getmem直接映射到异步 DMA 引擎和 NIC绕过 GPU 流式多处理器SM以节省计算资源。其次我们支持动态形状。为支持 MoE 等动态模型架构我们允许数据块大小、输入大小和输出大小在运行时动态计算。例如在 AllToAll 操作中基于 token 路由结果可以在运行时动态确定在 ranks 之间传输的块大小。任务级接口在最高级别Ditron 将整个分布式工作负载视为一个有向无环图DAG的任务。此抽象允许编译器将通信和计算内核融合到单个 MegaKernel超内核中。通过全局管理任务依赖关系和调度Ditron 消除了重复内核启动的开销并实现了内核在硬件上的持久驻留确保通信结构和计算单元持续繁忙。与先前需要 CUDA 或 C 编程的 MegaKernel 工作 [25, 29] 不同Ditron 允许用户将原生 Triton 内核注册为任务然后在编译时通过软件维护的记分板自动调度这些内核。最后这些 Triton 内核被转换为融合的 MegaKernel。我们在附录 D.4 中展示了如何将 Triton 内核注册为任务以及最终生成的 MegaKernel。3.2 中端计算-通信交织Swizzling中端负责将高级计算-通信算法转换为重叠版本并应用系统感知的优化。Ditron 中最关键的优化是分布式交织distributed swizzling。图4 在 8× H800 GPU 上的单工作负载评估。工作负载包括 AllGatherGEMM (AG-GEMM)、GEMMReduceScatter (GEMM-RS)、GEMMAllReduce (GEMM-AR)、AllGatherMoE (AG-MoE)、MoEAllReduce (MoE-AR)。由于 MoE 的加速比跨度较大故采用对数刻度具体数值列于附录 C。在分布式环境中访问远程内存的延迟比本地 HBM 高出数个数量级。标准的顺序执行通常会在等待数据时使计算单元空闲。Ditron 通过重新排序分块的执行来解决这个问题。我们将此形式化为编译器处理的两种不同模式收集模式Gather Mode对于计算依赖于远程数据的操作例如AllGatherGEMM编译器会尽可能早地调度远程数据获取请求有效地将本地 HBM 视为远程内存的缓存。分散模式Scatter Mode对于必须将结果发送出去的操作例如GEMMReduceScatter编译器会优先计算目标为最远远程节点的分块确保在数据可用时立即启动长途通信。交织逻辑在编译器中封装为一个无状态的、兼容 JIT 的工具。利用 Triton 的 JIT 基础设施和 Ditron 的三级接口我们将交织作为一个可插拔的原语暴露出来可以灵活地注入到任何内核中。该转换定义为我们在附录 A 中提供了详细的实现。为了说明其有效性图 3 可视化了 GEMMReduceScatter 工作负载的执行流程。在朴素调度无交织中所有 ranks 从索引 0 开始顺序处理分块。由于第 0 个分块所需的数据通常驻留在 Rank 0 上所有其他 ranks 在等待数据传输时会遭遇依赖停滞。相反Ditron 应用了 rank 感知的偏移量具体地从rank_id mod local_world_size开始执行。这确保了每个 rank 优先处理本地数据可用或即将到达的分块从而有效地将 GEMM 计算与 ReduceScatter 通信重叠。对于非完美形状的交织更具挑战性我们在附录 A 中进行了讨论。3.3 后端原语与代码生成原语为实现硬件可移植性设计原则 3Ditron 将后端特定的通信功能抽象为一组符合 OpenSHMEM 标准的硬件无关原语。在代码生成阶段Ditron 的分布式中间表示IR被降级为 LLVM IR。我们利用 LLVM 的 CallExtern 能力链接到供应商特定的通信库。这些原语分为三类分布式原语、SIMT 原语和 SHMEM 设备原语。分布式原语用于生成分布式信号控制的低级代码。SIMT对于训练我们分析了 TP、SP 和 EP 内核在 8 到 128 个 GPU 上的弱扩展和强扩展性能。最后我们展示了对 AMD 和 PCIe GPU 的支持。图5 Qwen3-32B 的模块级 TP 评估Prefill 和 Decode Attention 的结果显示与使用 AllGather 和 ReduceScatter 相比使用 AllReduce 的 Ditron 相比 CuBLASNCCL 获得了更好的加速比。对于 MLP结果表明对于大形状AllGather 和 ReduceScatter 的组合性能更好而对于小形状AllReduce 更优。图6 Qwen3-32B 和 LLaMA3-70B 在 8×H8008×H800 上的端到端 TP 评估结果Ditron 在大批量大小和大模型上相比 vLLM 展现出优势。4 评估我们在多种并行配置下对 Ditron 进行了评估。对于推理节点内张量并行我们对 5 种 GEMM/MoE 工作负载、Attention/FFN 模块以及端到端的 Qwen3-32B/LLaMA3-70B 模型进行了基准测试。4.1 推理评估单工作负载评估我们在 8×H800 GPU 上评估了 5 种不同的工作负载共 26 种配置每种工作负载具有多个来自真实世界模型如 LLaMA [8]、Mixtral [12]、GPT [21]、Qwen [30] 和 DeepSeek [6]的形状配置。详细的形状配置见附录 B.2。我们的基线包括 CuBLAS [18]NCCL [20]非重叠、TileLink [33]、FLUX [4] 和 COMET [31]结果如图 4 所示。总体而言对于 AG-GEMMDitron 相对于 CuBLASNCCL 的几何平均加速比为 1.43×1.43×相对于 TileLink 为 1.13×相对于 FLUX 为 1.09×。对于 GEMM-RSDitron 相对于 CuBLASNCCL 的几何平均加速比为 1.27×相对于 TileLink 为 1.02×相对于 FLUX 为 1.30×。TileLink 和 FLUX 不支持 AllReduce。对于 GEMM-AR相对于 CuBLASNCCL 的几何平均加速比为 1.32×。AG-GEMM 和 GEMM-RS 的加速主要来自通信的重叠而非更快的 GEMM实际上我们使用的是 Triton 的 GEMM其速度略慢于 CuBLAS 和 FLUX对于大输入形状近 87.5% 的通信延迟被计算隐藏。GEMM-AR 的加速来自使用 Ditron 实现的更快的 AllReduce它通过 NVLink Sharp 提供的多内存归约/广播支持单次one-shot和两次two-shot算法。对于 AG-MoEDitron 相对于 CuBLASNCCL 的几何平均加速比为 19.18×19.18×相对于 TileLink 为 1.89×相对于 COMET 为 1.06×。对于 MoE-AR相对于 CuBLASNCCL 的平均加速比为 13.89×。模块级评估我们将上述工作负载整合到 Attention 模块和 FFN 模块中。Attention 模块中的 QKV 投影和输出投影被 Ditron 内核替换。Prefill 和 Decode 性能如图 5 所示。Prefill 输入长度范围从 512 到 16kDecode 输出长度范围从 128 到 4k。两条线表示使用 AllGatherReduceScatter 与使用 AllReduce 相比相对于 CuBLASNCCL 加速比的差异。结果表明对于 Attention 模块Prefill 和 Decode 都更倾向于 AllReduce 通信因为 GEMM 的归约维度Attention 中的头维度大小较小128使得 GEMM 延迟不足以隐藏 AllGather 或 ReduceScatter 的通信延迟。Ditron 的 Attention 模块使用 AllReduce相对于 CuBLASNCCL 的几何平均加速比Prefill 为 1.12×Decode 为 1.26×。另一方面FFN 模块中的 GEMM 使用较大的中间维度大小在给定足够输入 token大批量或序列长度时GEMM 的延迟可以隐藏通信延迟因此对于大于 2k 的序列长度重叠 AllGather 和 ReduceScatter 优于使用 AllReduce。对于 128k tokensDitron 使用 AllReduce 的加速比为 1.17×使用 AllGatherReduceScatter 的加速比为 1.27×。图7 在 Hopper (96GB HBM) GPU 集群上扩展 TP 和 SP 工作负载结果以相对于 CuBLASNCCL 的加速比呈现形状来自各种真实世界的 LLM详见附录 B.2。Ditron 在多达 16 个 GPU 上为 AG-GEMM 保持了加速比。对于 32 个 GPU仅对大型 LLM 的形状才能实现加速比。GEMM-RS 的加速比在 8-32 个 GPU 上保持一致。GEMM-A2A 由于其可扩展的 AllToAll 而展现出最佳的弱扩展性。图8 在 8× H800 上使用 Ditron 的分布式 MegaKernel 对 Qwen3-8B、Qwen3-32B、LLaMA-70B 的推理延迟。端到端评估我们使用端到端模型 LLaMA3-70B [8] 和 Qwen3-32B [30] 评估 Ditron。我们将 Ditron 集成到 vLLM [13] 中并比较启用/禁用 Ditron 的性能。结果如图 6 所示。未启用 Ditron 的 vLLM [13] 仍然是一个强大的基线因为 vLLM 原生采用了专家设计的高效 AllReduce 内核。对于小于 128 的批量大小vLLM 在吞吐量上略优于 Ditron。但对于大于 128 的批量大小Ditron 相比 vLLM 实现了 5%−30% 的加速。特别是对于批量大小 512我们在图 6 中展示了 LLaMA3-70B 和 Qwen3-32B 在不同输入长度和输出长度下的详细加速比。这相当于 LLaMA3-70B 实现了 12k tokens/s 的吞吐量Qwen3-32B 实现了 17k tokens/s 的吞吐量。分布式 MegaKernel 评估对于单批次推理硬件资源通常未得到充分利用产生 MegaKernel 的任务级调度可以消除内核启动开销增加 SM 活动并提高端到端性能。我们将 Ditron 的分布式 MegaKernel 与 PyTorchEager 模式和 CUDAGraph 模式、DitronCUDAGraph 和 Mirage [29] 进行比较如图 8 所示。几何平均加速比相对于 Torch Eager 为 6.28×相对于 Mirage 为 1.73×相对于 TorchCUDAGraph 为 1.33×相对于 DitronCUDAGraph 为 1.11×相对于 vLLM 为 1.10×。我们还在附录 D.4 中添加了 MegaKernel 代码示例。4.2 训练评估对于训练我们使用 DITRON 评估了 TPAllGatherGEMM, GEMMReduceScatter、SPGEMMAllToAll和 EPMoE Dispatch 和 MoE Combine的强扩展和弱扩展性能。TP 和 SP 的结果如图 7 所示EP 的结果如图 9 所示。详细的形状配置见附录 B.2。更多结果在不同 GPU 上见附录 C。TP 工作负载在节点间产生大量通信网络通信的低带宽使得扩展效果不佳。因此我们仅在 8-32 个 GPU 上观察到 AllGatherGEMM 和 GEMMReduceScatter 的加速比强扩展总 token 数为 32768范围从 0.80× 到 1.71×相对于 CuBLASNCCL。对于加速比低于 1 的情况主要原因是分片后每个 GPU 上的 GEMM 变得太小无法隐藏通信延迟。作为对比GEMMAllToAll (GEMM-A2A) 在从 8 个 GPU 扩展到 128 个 GPU 的弱扩展每个 rank 的序列长度保持不变中提供了一致的加速比此时 GEMM 的形状保持不变。至于 EP我们设置每个 rank 8192 个 tokentopk 为 8隐藏维度为 7168dispatch 和 combine 性能在强扩展全局总共 512 个专家和弱扩展每个 rank 8 个专家下保持相似。相对于 PyTorchNCCL 实现的加速比范围从 1.04× 到 4.70×。图9 使用 DITRON 扩展 EP dispatch 和 combine。4.3 对其他平台的支持由于我们分布式 IR 的灵活性和所依赖的 Triton 编译器的兼容性DITRON 可以移植到更多硬件平台。目前我们已成功支持 AMD GPU 和 PCIe GPU。我们将初步结果放在附录 C 中。总体而言在 AMD GPU 上相对于 RocmBLASRCCL 的加速比范围为 2%−38%在 PCIe GPU 上相对于 CuBLASNCCL 的平均加速比为 8.33×。在从数十亿到数千亿参数不等的模型训练任务中不同层都得到了 DITRON 的加速。对于 Attention 模块我们采用 SP Attention。我们重叠了 GEMM 和 AllToAll 通信。与原始的 Megatron [16] 实现相比Attention 投影部分的加速比超过 20%。对于 MoE 模块我们重叠了 Dispatch、GroupedGEMM 和 Combine 操作。与原生 Megatron 实现相比端到端性能提升达到 10%。即使与高度调优的手写 CUDA 重叠实现FLUX [4]相比我们也达到了相当的性能同时将代码长度减少了一个数量级以上并将开发周期从数月缩短至数天。对于优化器我们支持了 Muon 优化器 [15] 的优化优化步骤加速比超过 20%。对于流水线并行我们实现了高效的 PP 通信内核。它可以在节点内仅使用 8 个 SMs 即可饱和带宽在节点间仅需 1 个 SM同时支持与其他层的灵活重叠。我们所有适配的内核与原生实现逐比特相同完全确保了训练的准确性和稳定性。推理集成。对于推理服务我们在云端部署了高性能 TP 推理服务。我们通过重叠和融合 AllReduce 与 GEMM 操作来加速 TP 推理支持包括 PCIe 和 NVLink GPU 在内的多种推理场景端到端性能加速约 20%。对于边缘推理例如机器人设备上的推理我们提供了高性能的 TP Attention 和 TP GEMM 实现端到端加速比超过 30%。5 结论分布式推理和训练变得必要需要更多研究人员能够编写分布式内核。本工作提出了 Ditron一个灵活且通用的分布式编译器具有用于重叠内核的核心级、设备级和任务级接口。研究人员可以使用 Ditron 为不同的并行策略编写高效的分布式内核其性能可与专家调优的内核相当或更优。