Ziji SHI(史子骥)

Ph.D.

National University of Singapore

Biography

I am an ML System Researcher intertested in building efficient training and inference systems. I graduated with a PhD degree from the National University of Singapore in 2026, where I was fortunate to be advised by Prof. Jialin Li and Wei Lin from Alibaba. I have previously worked at Sea AI Lab, Google XLA, Alibaba Platform for AI (PAI) team, and Apple Turi.

During my undergrad, I had the privilege (and fun) of spending four years with the NTU HPC club, where we won the Overall Championship at SC'17 and set the LINPACK World Record. I’ve hosted a MLSys seminar in Singapore.

Outside of my academic pursuits, I enjoy cooking (menu), hiking, and ham radio. My Erdős number is 5.

Download my resumé .

Interests

Large-Scale Machine Learning
Automatic Distributed Training
Distributed Systems

Education

Ph.D. Candidate, Computer Science, 2021 - 2026
National University of Singapore
BEng in Computer Science, 2015 - 2019
Nanyang Technological University
Visiting Student, Fall 2016
New York University

Publications

Ziji Shi, Chaoyi Ruan, Penghui Qi, Guangxing Huang, Xinyi Wan, Min Lin, Jialin Li

October 2025 In SOSP'25 SAA Workshop

Tetris: Efficient and Predictive KV Cache Offloading for Agentic and Reasoning Workloads

We present a predictive KV cache offloading mechainism that support ultra-long decoding phase in reasoning and agentic workloads.

Ziji Shi, Le Jiang, Ang Wang, Jie Zhang, Chencan Wu, Yong Li, Xiaokui Xiao, Wei Lin, Jialin Li

September 2025 In 54th International Conference on Parallel Processing (ICPP 2025)

TAPAS: Fast and Automatic Derivation of Tensor Parallel Strategies for Large Neural Networks

We present a framework that drastically speeds up the process of deriving the tensor parallel schedule for large neural networks by 160x.

Ziji Shi, Jialin Li, Yang You

November 2024 In Proceedings of the 2024 ACM Symposium on Cloud Computing

ParaGAN: A Scalable Distributed Training Framework for Generative Adversarial Networks

We present ParaGAN, a cloud-training framework for GAN, which demonstrates near optimal scaling performance over thousands of acclerators with system & training co-design.

Ziji Shi, Le Jiang, Jie Zhang, Xianyan Jia, Yong Li, Chencan Wu, Jialin Li, Wei Lin

June 2023 In ISCA'23 ASSYST Workshop

TAP: Efficient Derivation of Tensor Parallel Plans for Large Neural Networks

We present a framework that drastically speeds up the process of deriving the tensor parallel schedule for large neural networks.

Xianyan Jia, Le Jiang, Ang Wang, Wencong Xiao, Ziji Shi, Jie Zhang, Xinyuan Li, Langshi Chen, Yong Li, Zhen Zheng, Xiaoyong Liu, Wei Lin

July 2022 In USENIX ATC'22

Whale: Efficient Giant Model Training over Heterogeneous GPUs

Whale is a highly scalable and efficient distributed training framework for deep neural networks, introducing a hardware-aware parallel strategy and user-enabled model annotations for optimising large-scale model training, demonstrating its prowess by successfully training a multimodal model with over ten trillion parameters on a 512-GPU setup.

See all publications