718 citations · 881 across the 22 of their papers we have counts for
22 papers
CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution
Zihao Ye, Yingyi Huang, Hongyi Jin +11
GPU kernel agents and GPU programming languages have advanced separately, leaving expert kernels difficult to reproduce. Agents usually treat the compiler as a fixed black box and…
Tabular Numeric Stretch Transformation
Zihao Ye, Juyong Kim, Johnna Sundberg +2
Tabular data presents unique challenges for deep learning due to its heterogeneous nature, where numeric features exhibit diverse distributions, scales, and statistical properties.…
Event Tensor: A Unified Abstraction for Compiling Dynamic Megakernel
Hongyi Jin, Bohan Hou, Guanjie Wang +18
Modern GPU workloads, especially large language model (LLM) inference, suffer from kernel launch overheads and coarse synchronization that limit inter-kernel parallelism. Recent me…
AVO: Agentic Variation Operators for Autonomous Evolutionary Search
Terry Chen, Zhifan Ye, Bing Xu +20
Agentic Variation Operators (AVO) are a new family of evolutionary variation operators that replace the fixed mutation, crossover, and hand-designed heuristics of classical evoluti…
Axe: A Simple Unified Layout Abstraction for Machine Learning Compilers
Bohan Hou, Hongyi Jin, Guanjie Wang +7
Scaling modern deep learning workloads demands coordinated placement of data and compute across device meshes, memory hierarchies, and heterogeneous accelerators. We present Axe La…
FlashInfer-Bench: Building the Virtuous Cycle for AI-driven LLM Systems
Shanli Xing, Yiyan Zhai, Alexander Jiang +10
Recent advances show that large language models (LLMs) can act as autonomous agents capable of generating GPU kernels, but integrating these AI-generated kernels into real-world in…
MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs
Xinhao Cheng, Zhihao Zhang, Yu Zhou +17
We introduce Mirage Persistent Kernel (MPK), the first compiler and runtime system that automatically transforms multi-GPU model inference into a single high-performance mega-kerne…
Accelerating Large-Scale Reasoning Model Inference with Sparse Self-Speculative Decoding
Yilong Zhao, Jiaming Tang, Kan Zhu +10
Reasoning language models have demonstrated remarkable capabilities on challenging tasks by generating elaborate chain-of-thought (CoT) solutions. However, such lengthy generation…
TeleRAG: Efficient Retrieval-Augmented Generation Inference with Lookahead Retrieval
Chien-Yu Lin, Keisuke Kamahori, Yiyu Liu +11
Retrieval-augmented generation (RAG) extends large language models (LLMs) with external data sources to enhance factual correctness and domain coverage. Modern RAG pipelines rely o…
FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving
Zihao Ye, Lequn Chen, Ruihang Lai +8
Transformers, driven by attention mechanisms, form the foundation of large language models (LLMs). As these models scale up, efficient GPU attention kernels become essential for hi…
MagicPIG: LSH Sampling for Efficient LLM Generation
Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye +8
Large language models (LLMs) with long context windows have gained significant attention. However, the KV cache, stored to avoid re-computation, becomes a bottleneck. Various dynam…
NanoFlow: Towards Optimal Large Language Model Serving Throughput
Kan Zhu, Yufei Gao, Yilong Zhao +13
Large Language Models (LLMs) have resulted in a surging demand for planet-scale serving systems, where tens of thousands of GPUs continuously serve hundreds of millions of users. C…
vMCU: Coordinated Memory Management and Kernel Optimization for DNN Inference on MCUs
Size Zheng, Renze Chen, Meng Li +3
IoT devices based on microcontroller units (MCU) provide ultra-low power consumption and ubiquitous computation for near-sensor deep learning models (DNN). However, the memory of M…
Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
Ruihang Lai, Junru Shao, Siyuan Feng +16
Dynamic shape computations have become critical in modern machine learning workloads, especially in emerging large language models. The success of these models has driven the deman…
Punica: Multi-Tenant LoRA Serving
Lequn Chen, Zihao Ye, Yongji Wu +3
Low-rank adaptation (LoRA) has become an important and popular method to adapt pre-trained models to specific domains. We present Punica, a system to serve multiple LoRA models in…
Atom: Low-bit Quantization for Efficient and Accurate LLM Serving
Yilong Zhao, Chien-Yu Lin, Kan Zhu +7
The growing demand for Large Language Models (LLMs) in applications such as content generation, intelligent chatbots, and sentiment analysis poses considerable challenges for LLM s…
SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning
Zihao Ye, Ruihang Lai, Junru Shao +2
Sparse tensors are rapidly becoming critical components of modern deep learning workloads. However, developing high-performance sparse operators can be difficult and tedious, and e…
TensorIR: An Abstraction for Automatic Tensorized Program Optimization
Siyuan Feng, Bohan Hou, Hongyi Jin +8
Deploying deep learning models on various devices has become an important topic. The wave of hardware specialization brings a diverse set of acceleration primitives for multi-dimen…
FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems
Yuwei Hu, Zihao Ye, Minjie Wang +6
Graph neural networks (GNNs) are gaining increasing popularity as a promising approach to machine learning on graphs. Unlike traditional graph workloads where each vertex/edge is a…
DGL-KE: Training Knowledge Graph Embeddings at Scale
Da Zheng, Xiang Song, Chao Ma +6
Knowledge graphs have emerged as a key abstraction for organizing information in diverse domains and their embeddings are increasingly used to harness their information in various…
BP-Transformer: Modelling Long-Range Context via Binary Partitioning
Zihao Ye, Qipeng Guo, Quan Gan +2
The Transformer model is widely successful on many natural language processing tasks. However, the quadratic complexity of self-attention limit its application on long text. In thi…
Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks
Minjie Wang, Da Zheng, Zihao Ye +12
Advancing research in the emerging field of deep graph learning requires new tools to support tensor computation over graphs. In this paper, we present the design principles and im…