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

TL;DR

DITRON is a distributed tensor compiler that introduces a hierarchical abstraction to optimize large language model performance across heterogeneous hardware, achieving significant speedups and enterprise deployment.

Contribution

It presents a novel multi-level tiling compiler with hierarchical abstraction supporting diverse parallelism for distributed tensor programs.

Findings

Achieves 6-30% speedup on kernels and 5-30% on inference.

Demonstrates portability across NVIDIA and AMD hardware.

Deployed at enterprise scale, saving GPU hours and improving training and inference efficiency.

Abstract

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. \ours{} has been deployed at the enterprise level for both training and inference. It achieves an MFU improvement of over 10\% in training tasks, saving approximately 500,000 GPU hours of training cost per month. For inference tasks, it delivers an end-to-end gain of over 20\% and has been applied to cloud service inference and edge inference scenarios.