About

Hi! I am a second-year PhD student in Computer Science at National University of Singapore advised by Yang You, where I also completed my master’s studies. I obtained my bachelor’s degree in CS & EE from Huazhong University of Science and Technology. Previously, I collaborated at Pika with Chenlin Meng and interned at Colossal-AI with Jiarui Fang.

I am actively looking for summer research internship in 2025. Please feel free to reach out if there are any opportunities available.

Research

My current research mainly focuses on efficient and scalable machine learning through algorithm, parallelism, scheduling and compiler optimization, recently with a primary emphasis on efficient video generation.

I am always happy to chat about interesting research ideas, and looking for academic collaborations and interns. Please drop me an email if you are interested in collaborating with me.

Selected Publications (all)

Efficient Video Generation

• Real-Time Video Generation with Pyramid Attention Broadcast
  | ICLR 2025 | first author | paper | code | blog |

• Training Variable Sequences with Data-Centric Parallel
  | arXiv | co-first author | code | blog |

• DSP: Dynamic Sequence Parallelism for Multi-Dimensional Transformers
  | arXiv | first author | paper | code |

Efficient Memory Cost

• AutoChunk: Automated Activation Chunk for Memory-Efficient Long Sequence Inference
  | ICLR 2024 | first author | paper | code |

• HeteGen: Heterogeneous Parallel Inference for Large Language Models on Resource-Constrained Devices
  | MLSys 2024 | first author | paper |

Efficient AI for Science

• FastFold: Optimizing AlphaFold Training and Inference on GPU Clusters
  | PPoPP 2024 | second author | paper | code |

Open-Source Projects

• VideoSys: An Easy and Efficient System for Video Generation
  Project lead.
• Colossal-AI: Making large AI models cheaper, faster and more accessible
  Core contributor (ranked 5th by 2024) with 38k star.

Experience

• Pika
  Host: Chenlin Meng
  Efficient video system.

• Colossal-AI
  Research Engineer Intern
  Host: Jiarui Fang and Shenggui Li
  Core contributor (ranked 5th by 2024) with 38k star, and responsible for features including ZeRO, TP, PP, compiler, and MoE.
  2022.07 - 2023.12

Education

• National University of Singapore
  Ph.D. in Computer Science
  2024.01 - now
• National University of Singapore
  M.S. in Computer Science
  2022.08 - 2023.12
• Huazhong University of Science and Technology
  B.S. in Computer Science & Electronic Information
  2018.09 - 2022.06