arXiv Machine Learning

One Megakernel is Not Enough: Predicting Where to Cut for Streaming Inference

The paper investigates how to optimally partition streaming inference models into megakernels and conventional graph components. It introduces a framework that predicts the best megakernel composition without full end‑to‑end runs, showing that multiple megakernels or a fully conventional graph often outperform a single megakernel across various models and batch sizes. The study defines four model categories based on megakernel stability and demonstrates significant latency and throughput gains on an H100 NVL GPU using a three‑megakernel design.

arXiv AI
Jun 30

KernelSight-LM: A Kernel-Level LLM Inference Simulator

arXiv:2606. 28565v1 Announce Type: cross Abstract: As large language models (LLMs) move into production serving, practitioners must rapidly evaluate inference performance across diverse hardware, models, and serving parameters to meet cost and latency targets.

By Xiteng Yao, Taeho Kim, Hengzhi Pei, Xinle Liu, Kyle Ulrich, Leonard Lausen, Ashish Khetan, Xiang Song, George Karypis, Martin Herbordt
arXiv Machine Learning
Sep 14

Dissecting GPU Utilization for LLM Inference on Nvidia Hopper

The paper investigates how a single SM utilization metric can misrepresent the true workload of large language model (LLM) inference on Nvidia Hopper GPUs. By profiling vLLM with FlashAttention‑3 and cuBLASLt on an H100 NVL across various phases (cold prefill, warm prefill, and decode) and varying sequence length and batch size, the authors replace the single utilization figure with eight detailed counter‑validated views. These views, tied to specific Nsight Compute counters or formulas, reveal how factors such as fragment fill, occupancy limits, stall signatures, wave quantization, and kernel selection create utilization gaps across four production models and six per‑layer kernel roles.

By Mohammad Siavashi, Gerald Q. Maguire Jr., Dejan Kostic, Marco Chiesa
arXiv AI
Sep 15

mKernel: Fast Multi-GPU, Multi-Node Fused Kernels

arXiv:2609.13585v1 Announce Type: cross Abstract: Communication has become a bottleneck in distributed training and inference of large models. Overlapping communication with computation at the granul...

By Ziming Mao, Yihan Zhang, Shawn Wei Chew, Shuang Ma, Costin Raiciu, Yang Zhou, Scott Shenker, Ion Stoica
arXiv Machine Learning
Jun 11

MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs

arXiv:2512. 22219v2 Announce Type: replace-cross Abstract: 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-kernel.

By Xinhao Cheng, Zhihao Zhang, Yu Zhou, Jianan Ji, Jinchen Jiang, Zepeng Zhao, Ziruo Xiao, Zihao Ye, Yingyi Huang, Ruihang Lai, Hongyi Jin, Bohan Hou, Mengdi Wu, Yixin Dong, Anthony Yip, Zihao Ye, Songting Wang, Wenqin Yang, Xupeng Miao, Tianqi Chen, Zhihao Jia
arXiv AI
Jul 28

X-Stage: An Overlooked Pipeline Stage for Communication-Computation Overlap in DiT Inference

arXiv:2607. 23264v1 Announce Type: cross Abstract: Fine-grained, device-initiated communication lets persistent GPU kernels in distributed diffusion transformer (DiT) inference issue remote stores and overlap data movement with Tensor Core computation.

By Jianwen Xian, Zhiyuan Xu, Yuchen Li, Ziliang Lai, Kang He, Zhen Huang, Aichen Feng, Jinyan Chen, Yilin Zhang, Qinqin Chen, Chengru Song
arXiv AI
Oct 21

BitMTP: When Multi-Token Prediction Meets Low-Bit Large Language Models

BitMTP is an algorithm‑system co‑design framework that combines multi‑token prediction (MTP) with ultra‑low‑bit large language model backbones to accelerate decoding while maintaining quality. It introduces joint next‑token and MTP objectives, multiple decoding modes (confidence‑only, speculative‑only, hybrid), and an accepted‑prefix alignment mechanism that keeps causal key‑value boundaries and CUDA Graph replay states consistent. Experiments on NVIDIA H20, RTX 3090, and Jetson Orin show speedups ranging from 2.0× to 3.3× across various ultra‑low‑bit models without sacrificing task performance.

By Ning Zhang, Shihao Wang, Jinrui Zhang, Chaodong Xiao, Lei Zhang
arXiv AI
Jun 3

Fine-Tuning and Serving Gemma 4 31B on Google Cloud TPU: A Technical Comparison with GPU Baselines

arXiv:2605. 25645v2 Announce Type: replace-cross Abstract: We present the first end-to-end demonstration of fine-tuning and serving Google's Gemma 4 31B model on TPU hardware, providing an empirical comparison of TPU and GPU platforms for large language model adaptation.

By Jatin Kishnani, Mayank Goel, Amit Singh, Pulkit Agrawal, Sairanjan Mishra