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

Size Zheng, Xuegui Zheng, Hanshi Sun, Qi Hou, Wenlei Bao, Shiyu Li, Haojie Duanmu, Jin Fang, Chenli Xue, Chenhui Huang, Yuanqiang Liu, Renze Chen, Ningxin Zheng, Dongyang Wang, Li-Wen Chang, Liqiang Lu, Yun Liang, Jidong Zhai, Xin Liu
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:163996-164028, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-zheng26ai, title = {{DITRON}: Distributed Multi-level Tiling Compiler for Parallel Tensor Programs}, author = {Zheng, Size and Zheng, Xuegui and Sun, Hanshi and Hou, Qi and Bao, Wenlei and Li, Shiyu and Duanmu, Haojie and Fang, Jin and Xue, Chenli and Huang, Chenhui and Liu, Yuanqiang and Chen, Renze and Zheng, Ningxin and Wang, Dongyang and Chang, Li-Wen and Lu, Liqiang and Liang, Yun and Zhai, Jidong and Liu, Xin}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {163996--164028}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/zheng26ai/zheng26ai.pdf}, url = {https://proceedings.mlr.press/v306/zheng26ai.html}, 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.} }
Endnote
%0 Conference Paper %T DITRON: Distributed Multi-level Tiling Compiler for Parallel Tensor Programs %A Size Zheng %A Xuegui Zheng %A Hanshi Sun %A Qi Hou %A Wenlei Bao %A Shiyu Li %A Haojie Duanmu %A Jin Fang %A Chenli Xue %A Chenhui Huang %A Yuanqiang Liu %A Renze Chen %A Ningxin Zheng %A Dongyang Wang %A Li-Wen Chang %A Liqiang Lu %A Yun Liang %A Jidong Zhai %A Xin Liu %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-zheng26ai %I PMLR %P 163996--164028 %U https://proceedings.mlr.press/v306/zheng26ai.html %V 306 %X 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.
APA
Zheng, S., Zheng, X., Sun, H., Hou, Q., Bao, W., Li, S., Duanmu, H., Fang, J., Xue, C., Huang, C., Liu, Y., Chen, R., Zheng, N., Wang, D., Chang, L., Lu, L., Liang, Y., Zhai, J. & Liu, X.. (2026). DITRON: Distributed Multi-level Tiling Compiler for Parallel Tensor Programs. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:163996-164028 Available from https://proceedings.mlr.press/v306/zheng26ai.html.

Related Material