TauLang: A Language and Compiler for Space-Time Programming on Statically Scheduled Accelerators
Xuxiao Yang, Junchi Wu, Yuqi Peng, Feiyang Wu, Zhuoran Li, Wei Kong, Shuo Tao, Dongyi Fan 等 14 位作者
原文摘要
Statically scheduled accelerators such as Groq's Tensor Streaming Processor (TSP) achieve low-latency large language model inference by stripping dynamic synchronization hardware from the chip. With no caches, arbiters, or hardware schedulers, the compiler must fix the exact slice and issue cycle of every operation, so when an operation runs is a condition for correctness. Existing kernel languages cannot target such machines because they rely on hardware to resolve cycle timing and instead encode schedules implicitly in code structure. We present TauLang, a kernel language and compiler that makes time a first-class citizen of the program. The programmer schedules operations, loops, and composed kernels with parameterized explicit times, so exploring schedules means changing parameter values rather than rewriting code. TauLang validates explicit times algebraically in TauIR, its timed intermediate representation. It then solves tensor placement and cycle scheduling jointly in an integer linear program, using explicit times to group connected operations into timing clusters that prune the search space. On a cycleaccurate model of a TSP-like machine, joint search achieves a $1.40\times$ latency speedup across ten kernels over a space-first flow. A few explicit time pins keep latency within one cycle of the untimed schedules and scale joint search to 1,636 GEMM blocks, whereas untimed search finds no schedule at 32.
cs.PL · cs.AR