Jingwei Chris Zuo

Google

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

Google

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

• 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

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

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

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

• 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

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

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