Photo of Xuan Jiang

Xuan Jiang

I work on high performance computing (HPC) and post-training for large models. I am a Sr. HPC/AI Engineer at AWS Annapurna Labs, where I work on GPU communication for collective operations and mixture-of-experts (MoE) expert parallelism. My changes have shipped in NVSHMEM 3.7.0 and vLLM, and I co-led the open-source release of DeepEP v2 on AWS. I am also a Research Affiliate at MIT. Before that, I was a Software Engineer at Google.

I did my Ph.D. at UC Berkeley, advised by Prof. James Demmel (SLICE Lab) and Prof. Raja Sengupta (Cal Unmanned Lab), with Prof. Alexandre Bayen (BAIR) on my committee.

I have mentored nine students, now at MIT, UIUC, Yale, ETH Zurich, UCSD, CMU, Northeastern, Duke, and De Anza College.

Email: xuanj [at] mit [dot] edu

Research


My HPC work covers GPU communication for collective operations and mixture-of-experts (MoE) expert parallelism, in NCCL, aws-ofi-nccl, NVSHMEM, DeepEP, and NCCL EP on EFA, and simulation that scales across many GPUs. My post-training work covers efficient fine-tuning, MoE compression, reinforcement learning, and self-evolving agents.

Research Overview

HPC: GPU Communication for Collectives and Expert Parallelism: For DeepEP, I wrote the dispatch kernel and reviewed the combine kernel of a kernel pair that needs only weak signals, enabling hybrid mode on transports without ordered delivery. The EFA version is open source in amazon-contributing/DeepEP. With NVIDIA's NVSHMEM team, I improved the libfabric transport, and the changes shipped in NVSHMEM 3.7.0. With NVIDIA's NCCL team, I co-lead the work to run NCCL EP on EFA. In aws-ofi-nccl, the NCCL network plugin for EFA, I moved the GDRCopy signal path of GPU-initiated networking (GIN) to a worker thread and made receive requests go back to their pool only after that signal completes, and in NCCL I added NCCL_GIN_PROXY_POLL_BATCH to batch the GIN proxy's work. I also fixed host heap corruption with NCCL 2.31 and later in vLLM.

HPC: Parallel Simulation: LPSim (TR Part C'24) simulates city-scale traffic across multiple GPUs with graph partitioning and ghost-zone vehicle migration. Devastator (PADS'24) is a scalable parallel discrete event simulation framework for modern C++. My sole-author PADS'24 paper designs a time-driven simulation framework for large-scale traffic networks.

Post-Training: Efficient Fine-Tuning and MoE: SMT (ICLR'25) updates only the most significant sub-matrices and uses 67% less GPU memory than full fine-tuning. ExpertCondenser (COLM'26) keeps the knowledge of rarely used, long-tailed experts during MoE fine-tuning. Less is MoE (EMNLP'26) shows that capabilities in MoE models concentrate in a few FFN intermediate dimensions and compresses experts by Fisher importance.

Post-Training: Reinforcement Learning and Agents: RAST-MoE-RL (ICML'26) uses a self-attention mixture-of-experts encoder in deep reinforcement learning for ride-hailing matching. CORAL (COLM'26) runs long-lived agents that explore and collaborate through shared persistent memory. I am also on the reef team, which uses interaction records and feedback to update model weights or the agent's harness of prompts, rules, and skills. CloudAnoAgent (COLM'26 Workshop) detects anomalies in cloud systems with an LLM agent and a neuro-symbolic mechanism.

Selected Work

For the full list of publications, please see Google Scholar.


* indicates co-first author † indicates project lead ✉︎ indicates corresponding author

Demo from the CORAL repository (Apache-2.0)
CORAL: Towards Autonomous Multi-Agent Evolution for Open-Ended Discovery
Ao Qu*, Han Zheng*, Zijian Zhou*, Yihao Yan, Yihong Tang, Shao Yong Ong, Fenglu Hong, Kaichen Zhou, Chonghe Jiang, Minwei Kong, Jiacheng Zhu, Xuan Jiang, Sirui Li, Cathy Wu, Bryan Kian Hsiang Low, Jinhua Zhao, Paul Pu Liang
Conference on Language Modeling (COLM), 2026
Transformer encoder, mixture-of-experts layer, and actor and critic heads in the RAST-MoE-RL model
RAST-MoE-RL: A Regime-Aware Spatio-Temporal MoE Framework for Deep Reinforcement Learning in Ride-Hailing
Yuhan Tang*, Kangxin Cui*, Jung Ho Park*, Yibo Zhao*, Xuan Jiang†✉︎, Haoze He✉︎, Jiangbo Yu, Haris N. Koutsopoulos, Jinhua Zhao
International Conference on Machine Learning (ICML), 2026
Framework with trainable router biases and always-active condenser experts beside routed experts
Preserving Long-Tailed Expert Information in Mixture-of-Experts Tuning
Haoze He, Xingyuan Ding, Xuan Jiang, Xinkai Zou, Alex Cheng, Yibo Zhao, Juncheng Billy Li, Heather Miller
Conference on Language Modeling (COLM), 2026
Less is MoE: Trimming Experts in Domain-Specialist Language Models
Haoze He*, Xinkai Zou*, Xuan Jiang, Xingyuan Ding, Ao Qu, Juncheng Billy Li, Heather Miller
Conference on Empirical Methods in Natural Language Processing (EMNLP), 2026
CloudAnoAgent: Anomaly Detection for Cloud Sites via LLM Agent with Neuro-Symbolic Mechanism
Xinkai Zou, Xuan Jiang†✉︎, Ruikai Huang, Haoze He, Parv Kapoor, Jinhua Zhao
COLM 2026 Workshop on AI Measurement, 2026
ReactionBench data pipeline: reaction videos are collected, split into aligned reactor and stimulus clips, filtered for quality, and annotated with facial action units
ReactionBench: A Human-Anchored Benchmark for Modeling Reactions to Video
Yibo Zhao*, Ao Qu*, Xuan Jiang, Jiahao Xu, Keane Ong, Hang Jiang, Zhaofeng Wu, Dingyi Zhuang, Yihong Tang, Kaichen Zhou, Jinhua Zhao, Paul Pu Liang
NeurIPS 2026 Workshop on Evaluation of Interactive Agents (IAEval), 2026
Diagram comparing LoRA's low-rank update with SMT's update of selected weight blocks
Sparse Matrix in Large Language Model Fine-Tuning
Haoze He*, Juncheng Billy Li*, Xuan Jiang, Heather Miller
International Conference on Learning Representations (ICLR), 2025
Map © OpenStreetMap contributors © CARTO
Large Scale Multi-GPU Based Parallel Traffic Simulation for Accelerated Traffic Assignment and Propagation
Xuan Jiang, Raja Sengupta, James Demmel, Samuel Williams
Transportation Research Part C: Emerging Technologies, 2024
Plot of committed-to-executed event ratio against core count, from 128 to 32,768
Devastator: A Scalable Parallel Discrete Event Simulation Framework for Modern C++
John Bachan, Jianlan Ye, Xuan Jiang, Tan Nguyen, Mahesh Natarajan, Maximilian Bremer, Cy Chan
ACM SIGSIM Conference on Principles of Advanced Discrete Simulation (PADS), 2024
Framework diagram with routing, graph partition, and traffic propagation across GPUs
Designing a Time-Driven Simulation Framework for Large-Scale Traffic Networks
Xuan Jiang
ACM SIGSIM Conference on Principles of Advanced Discrete Simulation (PADS), 2024
[Show all papers][Show selected]

Other Papers

From Patchwork to Network: A Comprehensive Framework for Demand Analysis and Fleet Optimization of Urban Air Mobility
Xuan Jiang, Xuanyu Zhou, Yibo Zhao, Shangqing Cao, Haoze He, Jinhua Zhao, Mark Hansen, Raja Sengupta
IEEE Transactions on Intelligent Transportation Systems, 2026
Evaluating Multi-Modal UAM Impacts at Metro Region Scale: A San Francisco Bay Area Study
Xin Peng, Xuan Jiang, Vishwanath Bulusu, Emin Burak Onat, Raja Sengupta
IEEE Transactions on Intelligent Transportation Systems, 2026
Simulation-Based Optimization for Vertiport Location Selection: A Surrogate Model with Machine Learning Method
Xuan Jiang, Shangqing Cao, Baichuan Mo, Junzhe Cao, Hao Yang, Yuhan Tang, Mark Hansen, Jinhua Zhao, Raja Sengupta
Transportation Research Record, 2025
DRBO: A Regional Scale Simulator Calibration Framework Based on Day-to-Day Dynamic Routing and Bayesian Optimization
Xuan Jiang, Yibo Zhao, Chonghe Jiang, Junzhe Cao, Alexander Skabardonis, Alex Kurzhanskiy, Raja Sengupta
Smart Cities, 2025
The Urban Air Mobility Fleet Scheduling Problem with Nonlinear Charging Time
Shangqing Cao, Emin Burak Onat, Xuan Jiang, Raja Sengupta, Mark Hansen
SSRN preprint, 2025
A Simulation-Optimization Framework for Developing Wind-Resilient AAM Networks
Emin Burak Onat, Shangqing Cao, Raiyan Rizwan, Xuan Jiang, Mark Hansen, Raja Sengupta, Anjan Chakrabarty
International Conference on Research in Air Transportation (ICRAT), 2024, Best Paper Award
Simulating Integration of Urban Air Mobility into Existing Transportation Systems: Survey
Xuan Jiang, Yuhan Tang, Junzhe Cao, Vishwanath Bulusu, Hao Yang, Xin Peng, Yunhan Zheng, Jinhua Zhao, Raja Sengupta
Journal of Air Transportation, 2024
Fleet Size and Spill for UAM Operation under Uncertain Demand
Shangqing Cao, Xuan Jiang, Emin Burak Onat, Bo Zou, Mark Hansen, Raja Sengupta, Anjan Chakrabarty
International Conference on Research in Air Transportation (ICRAT), 2024
Robust Reinforcement Learning Strategies with Evolving Curriculum for Efficient Bus Operations in Smart Cities
Yuhan Tang, Ao Qu†, Xuan Jiang†, Baichuan Mo†, Shangqing Cao, Joseph Rodriguez, Haris N. Koutsopoulos, Cathy Wu, Jinhua Zhao
Smart Cities, 2024
Entropy-Based Dynamic Programming for Efficient Vehicle Parking
Jean-Luc Lupien, Abdullah Alhadlaq, Yuhan Tang, Yan Wu, Jiayu Joyce Chen, Yutan Long, Xuan Jiang
IEEE International Conference on Intelligent Transportation Systems (ITSC), 2024
Integrating the Traffic Science with Representation Learning for City-Wide Network Congestion Prediction
Wenqing Zheng, Hao Frank Yang, Jiarui Cai, Peihao Wang, Xuan Jiang, Simon Shaolei Du, Yinhai Wang, Zhangyang Wang
Information Fusion, 2023
An Automated Machine Learning (AutoML) Method for Driving Distraction Detection Based on Lane-Keeping Performance
Chen Chai, Juanwu Lu, Xuan Jiang, Xiupeng Shi, Zeng Zeng
INFORMS Annual Meeting, 2023
A Metrics-based Method for Evaluating Corridors for Urban Air Mobility Operations
Xuan Jiang, Xin Peng, Vishwanath Bulusu, Cristian Poliziani, Gano Chatterji, Raja Sengupta
IEEE International Smart Cities Conference (ISC2), 2022
Causality and Advanced Models in Trip Mode Prediction: Interest in Choosing Swissmetro
Huy Pham, Xuan Jiang, Cong Zhang
International Conference on Transportation and Development (ICTD), 2022
Quantifying the Resilience of the US Domestic Aviation Network During the COVID-19 Pandemic
Aleksandar Bauranov, Steven Parks, Xuan Jiang, Jasenka Rakas, Marta C. González
Frontiers in Built Environment, 2021
Making Sense of Electrical Vehicle Discussions Using Sentiment Analysis on Closely Related News and User Comments
Xuan Jiang, Josh Everts
International Conference on Transportation and Development (ICTD), 2021
Chemical and Rheology Evaluation on the Field Short-Term Aging of High Content Polymer Modified Asphalt
Weidong Huang, Chuanqi Yan, Xuan Jiang
Transportation Research Board Annual Meeting, 2019

Open Source


My pull request adds opt-in dispatch and combine kernels that need only weak signals, enabling hybrid mode on transports without ordered delivery. I wrote the dispatch kernel and reviewed the combine kernel. The EFA version is open source in amazon-contributing/DeepEP, and my awsome-distributed-ai recipe builds it on EFA clusters.

NVSHMEM My commits

I worked with the NVSHMEM team on the libfabric transport: RMA batching, signal-path deadlock and hang fixes, completion-queue draining, and pre-release testing. The changes shipped in NVSHMEM 3.7.0.

NCCL My commit

Added NCCL_GIN_PROXY_POLL_BATCH, which lets the proxy of GPU-initiated networking (GIN) drain its requests in batches.

aws-ofi-nccl Merged PRs

In the NCCL network plugin for EFA, I moved the GDRCopy signal path of GIN to a worker thread and made receive requests return to their pool only after that signal completes.

Fixes host heap corruption with NCCL 2.31 and later.

rdmatop PR #48 (merged)

Adds headless CSV recording of RDMA counters at sub-millisecond sampling intervals.

reef connects agent inference, feedback, learning, and versioned delivery so agents can continually improve. I'm on the team, with 25 PRs authored and 18 merged, including reef's Claude Code harness adapter and its CORAL test-time-training example.

CORAL is an open-source research system where autonomous coding agents explore and collaborate through shared persistent memory. I co-authored the COLM 2026 paper and authored 14 PRs, with 12 merged.

My simulator for city-scale traffic across multiple GPUs, using graph partitioning and ghost-zone vehicle migration. It handles networks with over 223,000 nodes in seconds on one GPU.

Talks


GPUs: From CUDA to the Gigawatt Rack
Slides
[10/2026]
UC Berkeley, guest lecture in Control and Information Management
DeepEPv2 Deep Dive: Unbalanced Scenario and Performance Simulator
[07/2026]
AWS Annapurna Labs
Real-World Examples of LTTng, EFA Traces, Flame Graphs, and Histograms
[04/2026]
AWS Annapurna Labs
AI-Assisted Large-Scale Code Changes
[06/2025]
Google
Sparse Matrix Tuning: Efficient Fine-Tuning of Large Language Models
Slides
[04/2025]
MIT Senior Researcher Seminar
The Intuition Behind Attention
[07/2024]
Google
Large Scale Multi-GPU Based Parallel Traffic Simulation for Accelerated Traffic Assignment and Propagation
[05/2023]
Lawrence Berkeley National Laboratory, HPC Technology Department

Mentoring


Students I have mentored, and where they are now: