Publications

A list of publications. Bold indicates myself; * denotes equal contribution.

2026

  1. OSDI
    VTC: DNN Compilation with Virtual Tensors for Data Movement Elimination
    Muyan Hu, Ahan Gupta, Jiachen Yuan, Vima Gupta, Taeksang Kim, Xin Xu, Janardhan Kulkarni, Ofer Dekel, Vikram Adve, and Charith Mendis
    USENIX Symposium on Operating Systems Design and Implementation, 2026
  2. ISC
    MegaFold: System-Level Optimizations for Accelerating Protein Structure Prediction Models
    Hoa La*, Ahan Gupta*, Alex Morehead, Jianlin Cheng, and Minjia Zhang
    International Supercomputing Conference, 2026
    Kernels adopted by AMD for OpenFold3
  3. ICLR
    AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism
    Ahan Gupta*, Zhihao Wang*, Neel Dani, Masahiro Tanaka, Olatunji Ruwase, and Minjia Zhang
    International Conference on Learning Representations, 2026
    PDF Adopted by DeepSpeed PyTorch Foundation Blog

2025

  1. SC
    X-MoE: Enabling Scalable Training for Emerging Mixture-of-Experts Architectures on HPC Platforms
    Yueming Yuan, Ahan Gupta, Jianping Li, Sajal Dash, Feiyi Wang, and Minjia Zhang
    International Conference for High Performance Computing, Networking, Storage, and Analysis, 2025
    PDF Best Student Paper
  2. OOPSLA
    SPLAT: A Framework for Optimised GPU Code-Generation for Sparse Regular Attention
    Ahan Gupta, Yueming Yuan, Devansh Jain, Yuhao Ge, David Aponte, Yanqi Zhou, and Charith Mendis
    ACM Conference on Object-Oriented Programming, Systems, Languages and Applications, 2025
    PDF

In submission

  1. Preprint
    A Self-Pruning Transformer: Extreme KV-Cache Compression with Universal Attention
    Davis Wertheimer*, Haochen Shen*, Ahan Gupta*, Derrick Liu, Yu Chin Fabian Lim, Mudhakar Srivatsa, Raghu K. Ganti, Minjia Zhang, and Naigang Wang
    Under review
  2. Preprint
    FLuRKA: Fast Fused Low-Rank & Kernel Attention
    Ahan Gupta, Hao Guo, Yueming Yuan, Yanqi Zhou, and Charith Mendis
    Under review