LLM-Guided Autoscheduling for Large-Scale Sparse Machine Learning

Rubens Lacouture, Genghan Zhang, Konstantin Hossfeld, Tian Zhao, Kunle Olukotun

Keywords: LLM-guided autoscheduling, Sparse ML compilation, Dataflow accelerators, FPGA, Graph neural networks, fusion

TL;DR: Using an LLM as an autoscheduler for a sparse deep learning compiler targeting dataflow accelerators

Abstract:

Optimizing sparse machine learning (ML) workloads requires navigating a vast schedule space. Two of the most critical aspects of that design space include which ops to fuse and which loop/dataflow order to use within each fused region.
We present Autosparse, an LLM-guided autoscheduler atop a fusion-capable sparse ML compiler that focuses on fusion grouping and \emph{legal dataflow order} selection. The compiler enumerates lawful orders per fused region and exposes a lightweight FLOPs/byte signal; the LLM proposes structured candidates (fusion sets and orders) that we validate and rank before codegen. With backend defaults for blocking and parallelism held fixed, case studies on GCN, GraphSAGE show consistent gains over unfused baselines and parity with hand-tuned/heuristic schedules. Coupling LLM reasoning with compiler legality and roofline-style signals efficiently explores sparse scheduling spaces with minimal human effort.

Submission Number: 40