8 papers
MKEvolve: A Modular Multi-Agent Framework for Kernel Code Generation
Jason Yoo, Rajarshi Saha, Shaowei Zhu +3
Despite rapid progress in LLM-based code generation, writing correct and performant kernels for hardware accelerators remains a key bottleneck in scaling modern ML workloads. We pr…
Directional Alignment Mitigates Reward Hacking in Reinforcement Learning for Language Models
Wenlong Deng, Jiaji Huang, Kaan Ozkara +4
Reward hacking arises when a model improves a proxy reward by exploiting shortcuts rather than solving the intended task. We study this failure mode through the geometry of reinfor…
Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs
Hongyi Liu, Jiaji Huang, Zhen Jia +2
Speculative decoding is widely used in accelerating large language model (LLM) inference. In this work, we focus on the online draft model selection problem in speculative decoding…
TritonRL: Training LLMs to Think and Code Triton Without Cheating
Jiin Woo, Shaowei Zhu, Allen Nie +3
The rapid evolution of Large Language Models (LLMs) has driven a growing demand for automated, high-performance system kernels to accelerate machine learning workloads. We introduc…
ProxSparse: Regularized Learning of Semi-Structured Sparsity Masks for Pretrained LLMs
Hongyi Liu, Rajarshi Saha, Zhen Jia +5
Large Language Models (LLMs) have demonstrated exceptional performance in natural language processing tasks, yet their massive size makes serving them inefficient and costly. Semi-…
Inference Optimization of Foundation Models on AI Accelerators
Youngsuk Park, Kailash Budhathoki, Liangfu Chen +7
Powerful foundation models, including large language models (LLMs), with Transformer architectures have ushered in a new era of Generative AI across various industries. Industry an…