Jingwei Zuo(左京伟)
I am a second year Ph.D. student in the Computer Science department at Rice University, advised by Dr. Yuke Wang. Prior to this, I obtained bachelor’s degrees in both science and engineering from Tsinghua University.
My research focuses on efficient machine learning systems.
Email: jingwei.zuo [at] rice [dot] edu
Research
My main research question is:
👉 How to compute more with less resources?
💡 How I develop such a research focus?
Nowadays, modern deep neural networks, represented by large language models (LLMs), have an enormous number of parameters and consume significant amounts of energy. Scaling up the model to achieve superior capabilities is important, whereas keeping the cost down is also important. The energy OpenAI’s ChatGPT uses each year to respond to the users’ requests could power 43,204 U.S. homes for the entire year.[1] It is an outrageous number, which consolidates my belief that we should make every endeavor to cut down the cost of AI models, thereby making the new technology accessible to everybody and making the earth a greener one.
News
05/18/2026 🎉 Will join Google Cloud as a PhD SWE Intern!
Publications
We identify two distinct types of attention heads in large language models: retrieval heads, which require access to the full KV cache to locate relevant tokens, and streaming heads, which only need a constant-size cache. DuoAttention exploits this asymmetry to significantly reduce memory usage and decoding latency for long-context inference, while preserving model accuracy.
AgentVerse is a versatile framework designed to facilitate the deployment of multiple LLM-based agents in various applications. It provides a structured pipeline for multi-agent collaboration and uncovers emergent social behaviors arising from agent interactions, offering insights into collective intelligence in LLM-based systems.
Open Source Contributions
Open-source JAX framework for high-performance LLM training on TPUs and GPUs. Implemented Streaming DiLoCo for low-communication training across bandwidth-constrained datacenters.
Reference implementation of DuoAttention, splitting attention heads into retrieval and streaming heads to cut long-context inference memory by 2.1x and decoding latency by 3.0x. Published at NeurIPS 2025.
Framework for deploying multiple LLM-based agents in collaboration, covering reasoning, tool use, and embodied AI. Published at ICLR 2024.
Experiences
- Implemented Streaming DiLoCo (Distributed Low-Communication Training) in MaxText, Google Cloud's open-source JAX training framework, heavily reducing cross-datacenter (DCN) bandwidth usage while keeping convergence identical to synchronous data-parallel training.
- Leveraged JAX, Pathways, and core Google infrastructure to build asynchronous, communication-overlapping parameter synchronization pipelines, enabling efficient training across isolated, bandwidth-constrained clusters.
- Profiled distributed ML workloads with XProf to resolve JAX compiler sub-optimalities and system bottlenecks.
- Ran large-scale pre-training convergence experiments comparing DiLoCo against data-parallel baselines on 7B dense and 30B MoE models across hundreds of TPU v5p chips.
- Presented at the Google PhD Summit 2026 and TPU B3 Week Summit at the Mountain View office.
Moonshot AI
Beijing, China2024.12-2025.05
AI Infrastructure Internship
LLM Inference Team
- Optimized LLM serving through speculative decoding, model parallelism, and effective GPU scheduling on vLLM/SGLang-like serving frameworks.
- Analyzed CPU overhead and GPU kernel time by profiling inference systems with Nsight Systems and the PyTorch profiler.
- Compared parallelism strategies at scale, such as data parallel + expert parallel vs. tensor parallel + expert parallel for large models across multiple nodes.
Carnegie Mellon University
Pittsburgh, PA, USA2024.07-10
Remote Research Internship
InfiniAI Lab, advised by Prof. Beidi Chen
- Accelerated long-context LLM inference using top-k sparse attention, supporting long-context generation with minimal latency.
- Applied approximate nearest neighbor search (ANNS) to retrieve high-attention key-value pairs, reducing GPU memory overhead.
Massachusetts Institute of Technology
Cambridge, MA, USA2023.10-2024.05
Research Internship
Han Lab, advised by Prof. Song Han
DuoAttention ···
- Co-designed DuoAttention, a framework reducing inference memory by 2.1x and decoding latency by 3.0x for Llama-2/3 and Mistral models.
- Devised a method applying full KV caching to retrieval heads while using a constant-length KV cache for streaming heads, enabling scalable long-context inference with negligible accuracy loss.
- Resulted in a NeurIPS 2025 poster publication.
Tsinghua University
Beijing, China2023.03-09
Research Internship
Natural Language Processing Lab (THUNLP), advised by Prof. Zhiyuan Liu
AgentVerse ···
- Co-developed AgentVerse, a framework enabling multi-agent collaboration for reasoning, tool use, and embodied AI.
- Validated the framework across diverse settings including reasoning, coding, and tool utilization.
- Open-sourced the code on GitHub, where it continues to be widely used, resulting in an ICLR 2024 poster publication.
Educations
Rice University
2025.09-Present
Doctoral Study in Computer Science and Information Systems
Tsinghua University
2021.09-2025.6
B.Eng. in Electrical Engineering
B.S. in Fundamental Sciences (Math & Physics)
Northeastern University
2023.09-12
Exchange student at College of Engineering
Selected on Dean's list
Academic Services
Service as a reviewer for the following conferences:
- NeurIPS 2025
To Learn More About Me
Ideals
I would love to witness a world where humans could obtain more convenience, harmony, and happiness. Undeniably, my current research interest is only one minute factor contributing to this grand (and probably quixote) ideal. But the thing is, I would not like my research to go against this prospect at any time and under any circumstance. I advocate for the open source community.
Other Experiences
I went to Northeastern University for a one-semester exchange program in 2023 Fall and had a gorgeous time there! I love traveling around and have been to Hong Kong, Macao, Japan, Singapore, Australia, the US and of course many places of interest in mainland China.
Fun Facts
When I get nervous, I like to scratch my hair😬. So next time you see me doing that in a debate, you know you've got me there.
Hobbies
- Sports: Tennis.
- Music: Sometimes I play piano. Classical music and pop music are both my favorites.
- Photography: At the crossing between nature and humanity. Portfolio: (Ig) chris.z_pics
- Reading: books about sociology, psychology, and history, also novels.
Contacts
Feel free to reach out to me by email! We may even have an in-person coffee-chat if we are in the same city! I am always glad to talk to someone else, because other’s talk often inspires me and my words may inspire others too:)
Currently in Houston, TX · drag to spin