DITRON: Distributed Tile-Level Compiler

• The paper introduces DITRON, a distributed, tile-level compiler that extends Triton’s single-device model to full clusters using a three-layer abstraction (Core, Device, Task).
• It integrates multi-level tiling with swizzling-based scheduling to optimize compute–communication overlap, thereby reducing latency in large language model workloads.
• Comprehensive evaluations demonstrate that DITRON outperforms hand-tuned CUDA+NCCL libraries on NVIDIA, AMD, and PCIe systems, achieving significant speedups in key operations.

Searching arXiv for the DITRON paper and a small set of directly mentioned related systems to ground citations. DITRON is a distributed, tile-level compiler for parallel tensor programs that extends the Triton tile programming model from a single device to a full cluster and introduces a hierarchical programming abstraction spanning Core, Device, and Task levels (Zheng et al., 2 May 2026). It is designed for LLM training and inference on heterogeneous distributed hardware, where communication can account for $20\%$–$80\%$ of end-to-end time and where the limiting factor is the mapping of tensor programs onto a non-uniform, multi-level memory and interconnect hierarchy. Its stated objective is to retain the programmability of Triton while achieving performance on par with or better than hand-tuned CUDA+NCCL libraries, with compute–communication overlap exposed as a first-class, compilable construct (Zheng et al., 2 May 2026).

1. Problem setting and design target

Modern distributed LLM execution is framed as a tension between two existing approaches. High-performance CUDA/NCCL-style libraries such as CuBLAS, CUTLASS, NCCL, FLUX, COMET, and DeepEP provide expert-tuned kernels and collectives, exploit hardware features such as Tensor Cores, NVLink/NVSwitch “sharp” reductions, and LL protocols, and fuse computation with communication. However, they are described as rigid and non-programmable, and extending them to new model architectures or retargeting them to new hardware can require months of low-level engineering (Zheng et al., 2 May 2026).

Tensor compilers and DSLs such as TVM, TensorIR, Triton, TileLang, Pallas, TileLink, CoCoNet, DistEinsum, and Iris provide flexible, programmable interfaces and automatic scheduling for single-device kernels or operator-level distributed computations. They excel at per-operator tiling, fusion, and exploitation of device-local memory hierarchies, but they typically do not model the full distributed memory hierarchy well. Inter-node communication is often treated as opaque library calls, intra-node fabric and on-switch capabilities are underexposed, and most systems operate at operator granularity rather than tile granularity across devices. DITRON is positioned precisely against this gap: a system that “feels like Triton,” exposes a multi-level tiling model matching the hierarchy from SMs to RDMA, and provides portable, hardware-agnostic primitives for NVIDIA and AMD backends (Zheng et al., 2 May 2026).

A recurrent misconception in this area is that a compiler for tensor programs is sufficient if it can tile within a single GPU and invoke collectives externally. DITRON’s formulation rejects that premise. Its central claim is that for LLMs combining tensor parallelism, sequence parallelism, expert parallelism, and pipeline/model parallelism, the compiler must treat communication structure and overlap policy as part of the kernel compilation problem rather than as an after-the-fact library composition problem.

2. Compiler stack and execution workflow

DITRON is organized as a three-layer compiler stack consisting of a front-end, mid-end, and back-end (Zheng et al., 2 May 2026). The front-end presents Triton-style programming with three tiling levels and an extended set of distributed primitives, including signals, remote memory, and swizzling. The mid-end performs “Swizzle” transformations and compute–communication overlap scheduling, mapping operator-level collectives into tile-level semantics and reordering tile execution in “gather” and “scatter” modes to hide latency. The back-end introduces a hardware-agnostic Distributed IR with OpenSHMEM semantics, lowers it to LLVM IR, and targets hardware-specific implementations through NVSHMEM for NVIDIA and rocSHMEM for AMD (Zheng et al., 2 May 2026).

The workflow is explicit. A kernel author writes a Triton-like kernel that includes both computation and distributed memory operations at tile granularity. Standard compute operations are lowered to Triton IR and Triton GPU IR, while distributed operations are lowered to Distributed IR. The mid-end then rewrites tile indices through a swizzle function of the form

$$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$

after which the back-end generates GPU assembly together with calls into NVSHMEM, rocSHMEM, or vendor libraries to implement remote loads, stores, and signals. Optionally, multiple task kernels are fused into a persistent MegaKernel that executes a DAG of tasks such as attention, MLP, and all-reduce on the GPUs (Zheng et al., 2 May 2026).

The compiler therefore does not merely lower a distributed operator into a sequence of independent kernels. It attempts to preserve tile-level coordination across computation, communication, and task orchestration, which is the basis for its overlap strategy.

3. Core–Device–Task hierarchy

The key conceptual contribution is a three-level abstraction aligned with the hardware hierarchy (Zheng et al., 2 May 2026).

Level	Represents	Principal mechanisms
Core	SMs/warps/threads; registers, shared memory/SRAM, L2, HBM	Triton tile semantics, tl.program_id, vectorized loads/stores, tl.dot(), distributed-aware primitives
Device	Inter-device movement within a node and across nodes	Chunks, DMA/RDMA, putmem, getmem, putmem_nbi, putmem_signal
Task	Model-level DAGs across the cluster	Task metadata, software Scoreboard, persistent MegaKernel

At the Core level, DITRON inherits Triton’s tile semantics: kernels operate on static-shaped tiles such as $128 \times 128$ blocks in GEMM, use Triton-style thread-block indexing, and extend the language with distributed-aware primitives including dl.wait, dl.notify, dl.consume_token, dl.rank(), dl.num_ranks(), and dl.symm_at(ptr, rank). Core-level code is intentionally close to conventional Triton kernels; the notable addition is explicit signal-based control that allows a consumer kernel to proceed only when communication or a prior producer stage has made the required tiles available (Zheng et al., 2 May 2026).

At the Device level, the abstraction shifts from tiles to chunks, which are coarse-grained blocks composed of multiple core-level tiles and are moved via DMA or RDMA rather than SM-driven loads and stores. This level corresponds to NVLink, xGMI, PCIe, and NIC-based transfers and supports dynamic shapes and runtime calculations of chunk sizes, which is especially relevant for MoE workloads where the number of tokens per expert per rank is data-dependent (Zheng et al., 2 May 2026).

At the Task level, the unit of orchestration is a kernel invocation or tile-wise kernel invocation such as “QKV projection,” “attention split,” “allreduce,” or “MLP FC1.” Tasks carry metadata including task_type, layer_id, task_id, tile_id_or_start, and dependency ranges into a global task_deps array. Readiness is tracked by a software Scoreboard, and multiple tasks can be fused into a persistent MegaKernel in which Scoreboard.wait_deps blocks until predecessors complete and task execution releases tiles to wake dependents (Zheng et al., 2 May 2026).

This hierarchy is not only descriptive. It is the organizing principle by which DITRON maps per-tile Tensor Core execution, inter-GPU chunk transfers, and cluster-scale model DAG scheduling into a single compilation framework.

4. Multi-level tiling, communication integration, and overlap

DITRON’s tiling strategy is explicitly multi-level. Core tiles map to registers, shared memory, L2, and HBM; device-level chunks are groups of tiles sized to amortize DMA and NIC overhead; task-level partitions correspond to larger segments of the model graph such as an attention block followed by an MLP or an entire pipeline stage (Zheng et al., 2 May 2026). For GEMM $C = A \times B$ with $A \in \mathbb{R}^{M \times K}$, $B \in \mathbb{R}^{K \times N}$, and $C \in \mathbb{R}^{M \times N}$, the paper describes tiling along $M$ and $N$ and uses rank-aware reindexing such as

$80\%$0

together with swizzled program IDs to control the order in which tiles are requested and computed (Zheng et al., 2 May 2026).

The central optimization is swizzling. In Gather mode, as in AllGather+GEMM, data required for local compute resides on remote ranks, so tiles are reordered such that requests for remote tiles are issued early and local HBM acts as a cache for remote memory. In Scatter mode, as in GEMM+ReduceScatter, the compiler prioritizes tiles whose outputs must go to the farthest ranks and begins computation with those tiles. The paper states that this swizzling handles both “perfect” tiling and non-perfect shapes in which tiles belong to multiple ranks for subsequent ReduceScatter or AllToAll (Zheng et al., 2 May 2026).

Communication is abstracted through OpenSHMEM-style primitives in Distributed IR, including rank(), num_ranks(), symm_at(ptr, rank), wait, notify, consume_token, putmem, getmem, putmem_nbi, putmem_signal, and barrier_all. Collectives are integrated into tiling rather than represented as opaque external operations. AllGather is implemented as tile- or chunk-level putmem and getmem with per-rank barriers; ReduceScatter is hierarchical, using NVLink or xGMI intra-node and ring or trees inter-node; AllReduce is implemented both as a standard ring and via NVLink Sharp multimem instructions multimem_ld_reduce_v4 and multimem_st_v4 (Zheng et al., 2 May 2026).

The overlap algorithm is producer–consumer based. In AllGather+GEMM, communication is the producer and GEMM is the consumer; in GEMM+ReduceScatter, GEMM is the producer and ReduceScatter is the consumer. The consumer uses dl.wait and dl.consume_token to block on signals and then consume the produced data, while the producer advances to further tiles. The implementation also includes low-latency protocols similar to NCCL’s LL, D2D copy fusion to avoid driver overhead and stragglers, and PCIe-aware software barriers for systems lacking strong atomic guarantees (Zheng et al., 2 May 2026).

Within this communication model, DITRON provides specialized kernels for Tensor Parallelism (ag_gemm, gemm_rs, gemm_allreduce), Sequence Parallelism (sp_ag_attn, ulysses_comm), Expert Parallelism / MoE (a2a_single_gemm, ag_moe, moe_ar), and Pipeline Parallelism, where PP communication kernels are reported to saturate bandwidth using few SMs and overlap with compute (Zheng et al., 2 May 2026).

5. Evaluation, portability, and deployment

The evaluation spans NVIDIA, AMD, and PCIe systems, with workloads including micro-benchmarks, module-level attention and MLP from Qwen3-32B, end-to-end inference in vLLM on Qwen3-32B and LLaMA3-70B, single-batch inference with distributed MegaKernel, and training workloads covering TP, SP, and EP on 8–128 GPUs (Zheng et al., 2 May 2026). On NVIDIA systems, the paper reports 8× H800 NVLink nodes with 200 GB/s uni-directional NVLink and 8 NICs per node at 50 GB/s inter-node, as well as Hopper 96GB HBM with NVLink at 450 GB/s and 4 NICs per node at 25 GB/s inter-node. AMD evaluation uses 8× AMD GPUs connected with xGMI, and a separate PCIe setting uses 8× PCIe GPUs (Zheng et al., 2 May 2026).

Setting	Reported result	Baseline
H800, AG-GEMM	Geometric mean $80\%$1	CuBLAS+NCCL
H800, GEMM-RS	$80\%$2	CuBLAS+NCCL
H800, GEMM-AR	$80\%$3	CuBLAS+NCCL
H800, AG-MoE	$80\%$4	CuBLAS+NCCL
vLLM inference	$80\%$5–$80\%$6 throughput improvement for batch size $80\%$7	native vLLM
MegaKernel inference	$80\%$8 vs PyTorch eager	PyTorch eager
AMD, AG-GEMM	Geometric mean $80\%$9	RocmBLAS+RCCL
PCIe overall	Geometric mean $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$0	CuBLAS+NCCL

On isolated kernels over 8 H800 GPUs, the paper reports geometric mean speedups of $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$1 for AG-GEMM versus CuBLAS+NCCL, $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$2 versus TileLink, and $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$3 versus FLUX; $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$4 for GEMM-RS versus CuBLAS+NCCL, $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$5 versus TileLink, and $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$6 versus FLUX; $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$7 for GEMM-AR versus CuBLAS+NCCL; $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$8 for AG-MoE versus CuBLAS+NCCL, $$new\_pid = swizzle\_func(old\_pid,\ shape\_to\_swizzle,\ rank,\ world\_size,\ block\_size)$$9 versus TileLink, and $128 \times 128$0 versus COMET; and $128 \times 128$1 for MoE-AR versus CuBLAS+NCCL. The paper characterizes these as corresponding to $128 \times 128$2 speedups over expert-tuned CUDA “overlap” libraries in many cases, while noting that the MoE baseline is severely communication-bound (Zheng et al., 2 May 2026).

For communication hiding, AG-GEMM and GEMM-RS on large shapes hide up to $128 \times 128$3 of communication latency under computation. At the module level on Qwen3-32B, attention prefill with AllReduce-based TP yields a $128 \times 128$4 speedup versus CuBLAS+NCCL, attention decode yields $128 \times 128$5, and FFN at long sequences such as 128k tokens shows $128 \times 128$6 for AllGather+ReduceScatter versus $128 \times 128$7 for AllReduce (Zheng et al., 2 May 2026).

In end-to-end vLLM serving on 8× H800 TP, native vLLM is reported to be slightly faster for batch size $128 \times 128$8, whereas DITRON yields $128 \times 128$9–$C = A \times B$0 throughput improvement for batch size $C = A \times B$1. At batch size 512, the reported throughput is approximately $C = A \times B$2k tokens/s for LLaMA3-70B and $C = A \times B$3k tokens/s for Qwen3-32B. For single-batch inference with distributed MegaKernel on 8× H800, the reported gains are $C = A \times B$4 versus PyTorch eager, $C = A \times B$5 versus Mirage, $C = A \times B$6 versus PyTorch + CUDA Graphs, and $C = A \times B$7 versus vLLM (Zheng et al., 2 May 2026).

Training results are more heterogeneous. For TP on 8–32 GPUs with fixed global tokens $C = A \times B$8, AG-GEMM and GEMM-RS show speedups from $C = A \times B$9 to $A \in \mathbb{R}^{M \times K}$0 versus CuBLAS+NCCL depending on shape and GPU count, and the paper explicitly notes that speedups below $A \in \mathbb{R}^{M \times K}$1 arise when GEMM becomes too small per GPU after sharding. For SP on 8–128 GPUs with constant per-rank sequence length, GEMM+AllToAll shows consistent speedups. For EP with MoE, under strong and weak scaling with 8192 tokens/rank, top-$A \in \mathbb{R}^{M \times K}$2, hidden size 7168, and global experts 512 or 8 per rank, the reported speedups are $A \in \mathbb{R}^{M \times K}$3–$A \in \mathbb{R}^{M \times K}$4 versus PyTorch+NCCL (Zheng et al., 2 May 2026).

Portability is a major evaluation axis. On AMD GPUs, the paper reports geometric mean speedups of $A \in \mathbb{R}^{M \times K}$5 for AG-GEMM and $A \in \mathbb{R}^{M \times K}$6 for GEMM-RS versus RocmBLAS+RCCL, along with $A \in \mathbb{R}^{M \times K}$7–$A \in \mathbb{R}^{M \times K}$8 speedups on Qwen3-32B attention prefill/decode and FFN, and overall gains of $A \in \mathbb{R}^{M \times K}$9–$B \in \mathbb{R}^{K \times N}$0 depending on workload. On PCIe GPUs, the reported gains are $B \in \mathbb{R}^{K \times N}$1 for GEMM+RS, $B \in \mathbb{R}^{K \times N}$2 for AllToAll MoE dispatch, $B \in \mathbb{R}^{K \times N}$3–$B \in \mathbb{R}^{K \times N}$4 for AllGather+MoE, $B \in \mathbb{R}^{K \times N}$5 for MoE+RS, $B \in \mathbb{R}^{K \times N}$6 for MoE+AR, and $B \in \mathbb{R}^{K \times N}$7 for int8 and $B \in \mathbb{R}^{K \times N}$8 for FP8 in A2A+GEMM, with an overall geometric mean speedup of $B \in \mathbb{R}^{K \times N}$9 versus CuBLAS+NCCL (Zheng et al., 2 May 2026). This suggests that the overlap and tiling strategy is not confined to high-bandwidth NVLink systems and remains useful on PCIe-limited topologies.

The deployment claims are unusually concrete. DITRON is reported as deployed at enterprise scale within ByteDance Seed and partners for both LLM training and inference. The paper states that it achieves an MFU improvement of over $C \in \mathbb{R}^{M \times N}$0 in production training workloads, saving approximately 500,000 GPU hours of training cost per month. For inference it reports more than $C \in \mathbb{R}^{M \times N}$1 end-to-end gain in cloud TP inference and more than $C \in \mathbb{R}^{M \times N}$2 speedup in edge inference for robotic devices. It also states that accelerated kernels are bitwise identical to native implementations, preserving training accuracy and stability (Zheng et al., 2 May 2026).

6. Relation to adjacent systems, limitations, and implications

DITRON is situated below system-level frameworks such as Megatron, DeepSpeed, ZeRO, GPipe, PipeDream, TorchTitan, and FSDP, as a kernel/compiler layer that can act as the implementation substrate for their parallel operators. Relative to Triton, TileLang, ThunderKittens, and FlashInfer, its distinguishing feature is that communication is integrated into a distributed, multi-level tiling model rather than treated as an external primitive. Relative to Pallas, TileLink, and Iris, the paper emphasizes its three-level hierarchy, a unified set of OpenSHMEM-like primitives for cross-vendor portability, and explicit task-level MegaKernel generation with software scoreboarding. Relative to CoCoNet, DistEinsum, Distal, and Distributed Halide, it is described as kernel-centric and tile-level rather than operator-granular. Relative to hand-optimized overlap and fusion libraries such as FLUX, COMET, AMD-fused, FlashOverlap, TokenWeave, and Centauri, it is presented as achieving similar or better performance while exposing a programmable interface (Zheng et al., 2 May 2026).

Several limitations are stated or implied. The optimized library currently covers key LLM building blocks rather than arbitrary operations; arbitrary operators may require manual kernel authoring. Topology-specific optimizations such as NVLink Sharp multimem depend on specific NVIDIA hardware, and irregular or exotic network topologies may require further tuning. Tile sizes and swizzling patterns are described as mostly crafted rather than auto-searched, so broader auto-tuning across multi-level tiling and distributed communication remains open. The paper also acknowledges scenarios in which strong scaling makes per-GPU GEMMs too small for communication to be fully hidden. Finally, although integration with vLLM and Megatron-like systems is demonstrated, broader frontend support such as PyTorch 2 compile, XLA, or JAX would require dedicated integration work (Zheng et al., 2 May 2026).

The broader implication is not simply that another distributed kernel library has been introduced. A more consequential reading is that DITRON treats distributed compute–communication overlap as a compiler transformation problem, embodied in a hierarchy of Core, Device, and Task abstractions, swizzling, and Distributed IR. This suggests a design point for future tensor compilers in which portability, tile-level programmability, and cluster-scale communication scheduling are co-designed rather than layered independently.