Jingwei Zuo (Chris)
/dzwaw Jeeng-way/
Kirkland, Washington
jingwei.zuo [at] rice [dot] edu
github
A carved seal reading 左京伟

discipline

Systems for AI, and AI for systems

position

2nd-year Ph.D., Computer Science, Rice University. Advised by Dr. Yuke Wang.

now

Student Researcher, Google Cloud AI Research
2026.09-Present

the question

How to compute more with less resources?

Two literatures grew up in different rooms. One asks what a language model can be made to do: plan, use tools, act as one of many. The other asks what it costs: memory, bandwidth, latency, watts. I work where they meet.

two directions

Efficient systems for AI

  • Sparse attention
  • Sparse architectures for deep learning
  • Efficient resource allocation and scheduling of deep-learning workloads

AI for efficient systems

  • Agents for kernel generation
  • Agents for building system infrastructure
  • Agent-automated LLM training

figure one

A mask,
sparsified

Every dot is a query–key pair a head could attend to. A causal mask is a triangle, because a token may only look backwards.

Scroll and watch a streaming head throw the middle away. It keeps a few attention sinks and a narrow window of recent tokens. The generated text does not notice, and the cache stops growing with the context.

Splitting heads into the ones that need everything and the ones that do not is DuoAttention: 2.1× less memory, 3.0× faster decoding.

pairs 2,080
cache O(n)
mask cut 0%

n = 64, drawn as a halftone. Four sink columns, a five-token window. At a 32k context, with the 64 sinks and 256-token window the paper uses, the same rule removes 98.0% of the mask.

publications

Two papers,
one argument

ICLR 2025

DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads

Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, Song Han

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.

papercodedemo
DuoAttention Illustration
DuoAttention Illustration
ICLR 2024

AgentVerse: Facilitating Multi-Agent Collaboration and Exploring Emergent Behaviors

Weize Chen*, Yusheng Su*, Jingwei Zuo, Cheng Yang, Chenfei Yuan, Chi-Min Chan, Heyang Yu, Yaxi Lu, Yi-Hsin Hung, Chen Qian, Yujia Qin, Xin Cong, Ruobing Xie, Zhiyuan Liu, Maosong Sun, Jie Zhou

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.

papercode
AgentVerse Illustration
AgentVerse Illustration

experience

Where the
hours went

2026.09-Present

Google

Student Researcher, Google Cloud AI Research · Managed by Linfeng Song

2026.05-2026.08
Seattle, WA, USA

Google

SWE Intern, JAX Training Experience Team, Core ML · Google Cloud

5 notes
  • 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.
MaxText #4759MaxText #4046MaxText #4829MaxText #4938
2024.12-2025.05
Beijing, China

Moonshot AI

AI Infrastructure Internship · LLM Inference Team

3 notes
  • 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.
2024.07-10
Pittsburgh, PA, USA

Carnegie Mellon University

Remote Research Internship · InfiniAI Lab, advised by Prof. Beidi Chen

2 notes
  • 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.
2023.10-2024.05
Cambridge, MA, USA

Massachusetts Institute of Technology

Research Internship · Han Lab, advised by Prof. Song Han

3 notes
  • 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.
2023.03-09
Beijing, China

Tsinghua University

Research Internship · Natural Language Processing Lab (THUNLP), advised by Prof. Zhiyuan Liu

3 notes
  • 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.

open source

Code that
left the desk

2026

MaxText

Google Cloud / AI-Hypercomputer · Contributor

Open-source JAX framework for high-performance LLM training on TPUs and GPUs. Implemented Streaming DiLoCo for low-communication training across bandwidth-constrained datacenters.

#4759#4046#4829#4938 (open)
2023-2024

DuoAttention

MIT Han Lab · Core contributor

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.

2023

AgentVerse

OpenBMB / THUNLP · Core contributor

Framework for deploying multiple LLM-based agents in collaboration, covering reasoning, tool use, and embodied AI. Published at ICLR 2024.

education & service

Rooms with
blackboards

2025.09-Present

Rice University

Doctoral Study in Computer Science and Information Systems

2021.09-2025.6

Tsinghua University

B.Eng. in Electrical Engineering

B.S. in Fundamental Sciences (Math & Physics)

2023.09-12

Northeastern University

Exchange student at College of Engineering

Selected on Dean's list

NeurIPS 2026

Reviewer

PPoPP 2026

Artifact evaluator

NeurIPS 2025

Reviewer

off duty

Four other
appetites

01

Sports

Tennis.

02

Music

Sometimes I play piano. Classical music and pop music are both my favorites.

03

Photography

At the crossing between nature and humanity. Portfolio: (Ig) chris.z_pics

04

Reading

books about sociology, psychology, and history, also novels.