DITRON: Distributed Multi-level Tiling Compiler for Parallel Tensor Programs

The scaling of large language models (LLMs) is currently bottlenecked by the rigidity of distributed programming. While high-performance libraries like CuBLAS and NCCL provide optimized primitives, they lack the flexibility required for rapidly evolving model architectures. Conversely, existing tensor compilers fail to address the complex memory hierarchy of distributed clusters effectively. To bridge this gap, we propose DITRON, a scalable tile-level compiler that democratizes high-performance distributed kernel development. DITRON introduces a novel hierarchical programming abstraction spanning Core, Device, and Task levels to map tensor programs efficiently onto heterogeneous distributed hardware. This abstraction allows DITRON to support diverse parallelism strategies while abstracting away the complexity of inter-node and intra-node communication. Evaluated across large-scale clusters, DITRON achieves performance parity with or exceeding expert-tuned CUDA libraries, delivering speedups of 6%–30% on isolated kernels and 5%–30% on end-to-end inference in vLLM. Furthermore, DITRON demonstrates strong portability, achieving significant speedups on both NVIDIA and AMD platforms.

Training and running today's large AI models requires thousands of specialized chips called GPUs working together as a team. But getting these chips to cooperate efficiently is surprisingly hard. The standard tools engineers use are like pre-built Lego sets, which are fast, but rigid: they only fit a few well-known model designs. When researchers invent a new kind of AI model, these tools often can't keep up, leaving expensive hardware sitting idle and wasting enormous amounts of energy. We built DITRON, a new kind of "translator" that takes the high-level description of an AI model and automatically turns it into highly efficient instructions for thousands of cooperating chips. The key idea is a layered design that thinks about the work at three scales at once, including what each individual processor does, what each chip does, and how the whole cluster coordinates, so the pieces always fit together smoothly, even for unusual or brand-new model designs. In real-world deployment at large scale, DITRON matches or beats the hand-tuned tools that took years to build, saves roughly half a million GPU-hours every month, and makes AI services noticeably faster and cheaper for end users whether running in massive data centers or on small edge devices.