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
Large language model (LLM)-based evolution is a promising approach for open-ended discovery, where progress requires sustained search and knowledge accumulation. Existing methods still rely heavily on fixed heuristics and hard-coded exploration rules, which limit the autonomy of LLM agents. We present CORAL, the first framework for autonomous multi-agent evolution on open-ended problems. CORAL replaces rigid control with long-running agents that explore, reflect, and collaborate through shared persistent memory, asynchronous multi-agent execution, and heartbeat-based interventions. It also provides practical safeguards, including isolated workspaces, evaluator separation, resource management, and agent session and health management. Evaluated on diverse mathematical, algorithmic, and systems optimization tasks, CORAL sets new state-of-the-art results on 10 tasks, achieving 3-10 times higher improvement rates with far fewer evaluations than fixed evolutionary search baselines across tasks. On Anthropic's kernel engineering task, four co-evolving agents improve the best known score from 1363 to 1103 cycles. Mechanistic analyses further show how these gains arise from knowledge reuse and multi-agent exploration and communication. Together, these results suggest that greater agent autonomy and multi-agent evolution can substantially improve open-ended discovery. Code is available at https://github.com/Human-Agent-Society/CORAL.
Ride-hailing platforms face the challenge of balancing passenger waiting times with overall system efficiency under highly uncertain supply-demand conditions. Adaptive delayed matching, which controls the holding intervals for batched sets of requests and vehicles, reveals an inherent trade-off between matching and pickup delays. The resulting environment with temporally varying request arrival patterns and dynamic congestion calls for more expressive networks with sufficient capacity to capture their non-stationarity. To address the limitations of existing methods that rely on shallow encoders that cannot capture dynamic supply-demand patterns and congestion effects, we introduce the Regime-Aware Spatio-Temporal Mixture-of-Experts (RAST-MoE) framework, which formalizes adaptive delayed matching as a regime-aware Markov Decision Process and equips RL agents with a self-attention MoE encoder. Instead of relying on a single monolithic network, our design allows different experts to specialize automatically in varying operational conditions, improving representation capacity while maintaining per-sample computation efficiency. Despite its modest size of only 12M parameters, our framework consistently outperforms strong baselines. On real-world Uber trajectory data from San Francisco, it reduces average matching delay by 10%, and pickup delay by 15%. In addition, it demonstrates robustness to unseen demand regimes, stable training behavior without reward hacking, and expert specialization to different regimes. This study shows the strength of MoE-enhanced RL for large-scale decision-making tasks with complex spatiotemporal dynamics.
Despite MoE models leading many benchmarks, supervised fine-tuning (SFT) for the MoE architectures remains difficult because its router layers are fragile. Methods such as DenseMixer and ESFT mitigate router collapse with dense mixing or auxiliary load-balancing losses, but these introduce noisy gradients that often degrade performance. In preliminary experiments, we systematically pruned experts and observed that while certain super experts are activated far more frequently, discarding less used experts still leads to notable performance degradation. This suggests that even rarely activated experts encode non-trivial knowledge useful for downstream tasks. Motivated by this, we propose an auxiliary-loss-free MoE SFT framework that combines bias-driven sparsification with always-active gated condenser experts. Rather than enforcing balanced activation across all experts, our method encourages task-relevant experts to remain active while pushing long-tailed experts toward inactivity. The condenser experts provide a persistent, learnable pathway that alleviates gradient starvation and facilitates consolidation of information that would otherwise remain fragmented across sparsely activated experts. Analysis further suggest that this design better preserves long-tailed expert information under sparse routing. Experiments on large-scale MoE models demonstrate that our approach outperforms state-of-the-art SFT baselines such as DenseMixer and ESFT, achieving average gain of 2.5%+ on both mathematical reasoning and commonsenseQA benchmarks.
Mixture-of-Experts (MoE) models achieve strong performance through conditional computation, but their large parameter footprint poses deployment challenges. Prior MoE compression approaches catastrophically fail when evaluated on general-purpose benchmarks beyond commonsense reasoning. We trace this failure to the granularity of compression: important capabilities are distributed across experts but concentrated in FFN sparse intermediate dimensions. To identify these dimensions, we use Fisher importance which outperforms activation-, router-score-, and magnitude-based alternatives, and identifies tiny sets of task-critical dimensions: in Qwen1.5-MoE, removing as few as 12 of 1.35M routed-FFN intermediate dimensions collapses GSM8K accuracy while largely preserving factual-knowledge performance. Building on this, we propose Fisher-MoE, which operates within FFN to remove intermediate dimensions ranked by Fisher importance. At the same 50% MoE compression ratio, Fisher-MoE preserves model capability, while reducing weight memory by ~45% and improving inference throughput by 21%. These findings suggest intermediate dimension granularity is an effective unit for both compression and ranking where capability concentrates in MoE models.
Anomaly detection in cloud sites remains a critical yet challenging task. Existing approaches that rely solely on metric data often suffer from high false positive rates (FPR) due to data imbalance between normal and anomalous events, leading to significant operational overhead for system reliance engineers. Recent advances in large language models (LLMs) offer new opportunities for integrating metrics with log data, enabling more accurate and interpretable anomaly detection. In this paper, we propose CloudAnoAgent, the first neuro-symbolic LLM-based agent for anomaly detection in cloud environments. CloudAnoAgent jointly processes structured metrics and textual log data in a unified pipeline, leveraging symbolic verification to validate detection hypotheses and generate structured anomaly reports. To support systematic evaluation, we introduce CloudAnoBench, the first benchmark that provides LLM-generated paired metrics and log data with fine-grained anomaly behavior annotations, filling a critical gap in existing datasets. Experimental results demonstrate that CloudAnoAgent improves anomaly classification accuracy by 46.36% and 36.67% on average and reduces the FPR by 36.67% and 33.89% on average over traditional baselines and LLM-only baseline, with a boost on anomaly type detection accuracy by 12.8% compared to vanilla LLM prompting. These results demonstrate the strengths of our approach in improving detection accuracy, reducing false positives, and enhancing interpretability, thereby supporting practical deployment in enterprise cloud environments.
Video benchmarks ask what happens in a video. Socially responsive systems also need to relate what a person watches to how that person responds: a laugh follows a visual gag; a pause matters only against the same viewer’s surrounding behavior. ReactionBench asks whether current multimodal models can rank, temporally match, and explain a viewer’s reaction against the stimulus it was watching. Across seven models, the answer is: not yet.
LoRA and its variants have become popular parameter-efficient fine-tuning (PEFT) methods due to their ability to avoid excessive computational costs. However, an accuracy gap often exists between PEFT methods and full fine-tuning (FT), and this gap has yet to be systematically studied. In this work, we introduce a method for selecting sparse sub-matrices that aim to minimize the performance gap between PEFT vs. full fine-tuning (FT) while also reducing both fine-tuning computational cost and memory cost. Our Sparse Matrix Tuning (SMT) method begins by identifying the most significant sub-matrices in the gradient update, updating only these blocks during the fine-tuning process. In our experiments, we demonstrate that SMT consistently surpasses other PEFT baseline (e.g. LoRA and DoRA) in fine-tuning popular large language models such as LLaMA across a broad spectrum of tasks, while reducing the GPU memory footprint by 67% compared to FT. We also examine how the performance of LoRA and DoRA tends to plateau and decline as the number of trainable parameters increases, in contrast, our SMT method does not suffer from such issue.
Traffic propagation simulation is crucial for urban planning, enabling congestion analysis, travel time estimation, and route optimization. Traditional micro-simulation frameworks are limited to main roads due to the complexity of urban mobility and large-scale data. We introduce the Large Scale Multi-GPU Parallel Computing based Regional Scale Traffic Simulation Framework (LPSim), a scalable tool that leverages GPU parallel computing to simulate extensive traffic networks with high fidelity and reduced computation time. LPSim performs millions of vehicle dynamics simulations simultaneously, outperforming CPU-based methods. It can complete simulations of 2.82 million trips in 6.28 minutes using a single GPU, and 9.01 million trips in 21.16 minutes on dual GPUs. LPSim is also tested on dual NVIDIA A100 GPUs, achieving simulations about 113 times faster than traditional CPU methods. This demonstrates its scalability and efficiency for large-scale applications, making LPSim a valuable resource for researchers and planners. Code: https://github.com/Xuan-1998/LPSim
[Show all papers][Show selected]
Other Papers
Urban Air Mobility (UAM) presents a transformative vision for metropolitan transportation, but its practical implementation is hindered by substantial infrastructure costs and operational complexities. We address these challenges by modeling a UAM network that leverages existing regional airports and operates with an optimized, heterogeneous fleet of aircraft. We introduce LPSim, a Large-Scale Parallel Simulation framework that utilizes multi-GPU computing to co-optimize UAM demand, fleet operations, and ground transportation interactions simultaneously. Our equilibrium search algorithm is extended to accurately forecast demand and determine the most efficient fleet composition. Applied to a case study of the San Francisco Bay Area, our results demonstrate that this UAM model can yield over 20 minutes' travel time savings for 230,000 selected trips. However, the analysis also reveals that system-wide success is critically dependent on seamless integration with ground access and dynamic scheduling.
We present Vertiport-informed Surrogate-Based Optimization with Machine Learning Surrogates (VinS), a novel framework for solving the vertiport location problem for urban air mobility operations. The primary focus of this work is on the optimization of vertiport locations to facilitate efficient urban air transportation. Our framework helps choose not only the optimal vertiport locations but also the optimal number of vertiports. We develop a simulation model to analyze the impacts of various vertiport location configurations on the efficiency of the transportation network. To expedite the simulation process, a surrogate model is trained using machine learning algorithms, effectively reducing the computational time for evaluating a given vertiport location configuration. With the machine learning surrogate models, we apply a genetic algorithm to find the optimal set of vertiport locations. An empirical study was performed in the San Francisco Bay Area, from which we found that given the optimal set of vertiport locations, we further reduced the total travel time in the entire transportation system in the Bay Area compared with the sampled sets by 0.05% (400 h) on average. We release our framework at: https://github.com/Xuan-1998/LPSim .
Environmental factors pose a significant challenge to the operational efficiency and safety of advanced air mobility (AAM) networks. This paper presents a simulation-optimization framework that dynamically integrates wind variability into AAM operations. We employ a nonlinear charging model within a multi-vertiport environment to optimize fleet size and scheduling. Our framework assesses the impact of wind on operational parameters, providing strategies to enhance the resilience of AAM ecosystems. The results demonstrate that wind conditions exert significant influence on fleet size even for short-distance flights, their impact on fleet size and energy requirements becomes more pronounced over longer distances. Efficient management of fleet size and charging policies, particularly for long-distance networks, is needed to accommodate the variability of wind conditions effectively.
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.
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.
Added NCCL_GIN_PROXY_POLL_BATCH, which lets the proxy of GPU-initiated networking (GIN) drain its requests in batches.
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.
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
Mentoring
Students I have mentored, and where they are now:
- Yuhan Tang (MIT, Ph.D. student)
- Yibo Zhao (UIUC, Ph.D. student working on coding agents)
- Lucas Huang (Yale University)
- Junzhe Cao (ETH Zurich, ESOP Scholar 2024)
- Xinkai Zou (UCSD, CS master's)
- Jiaying Li (CMU, CS master's)
- Luze Sun (Northeastern University, CS Ph.D. student)
- Jieying Zhang (Duke, ECE master's)
- Cy Xie (De Anza College, CS undergraduate)