Wave: A Symbolic Python DSL And Compiler for High-Performance Machine Learning

Harsh Menon, Oleksandr Zinenko, Gaurav Verma, Stanley Winata, Ivan Butygin, Nithin Meganathan, Sanket Pandit, William Hatch, Surya Jasper, Megan Kuo, Sahil Faizal, Ashay Rane, Aurore De Spirlet, Martin Paul Lücke

Proceedings of Machine Learning and Systems 8 (MLSys 2026) Conference

Modern ML models demand ever-greater compute, prompting hardware vendors to add specialized matrix cores to their GPUs. While these units unlock high throughput, they impose intricate programming models and addressing schemes that are difficult to manage by hand. This paper introduces Wave, a Python-embedded DSL for kernel authoring that automates these complex address computations and lets authors focus on core computation. In experiments, it matches or surpasses the performance of state-of-the-art kernel DSLs and libraries.