# ASTRA-sim 3.0: Next-Level Distributed Machine Learning Simulations via High-Fidelity GPU and Infrastructure Modeling

arXiv: https://arxiv.org/abs/2606.10440 · HTML: https://mlsanity.com/paper/2606.10440/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 51 references that are in the map.
- First posted on arXiv 9 Jun 2026.

## Builds on (most cited first, 18 of 18)

- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence](https://mlsanity.com/paper/2606.19348/) (2026)
- [ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale](https://mlsanity.com/paper/2303.14006/) (2023)
- [Blink: Fast and Generic Collectives for Distributed ML](https://mlsanity.com/paper/1910.04940/) (2019)
- [LLMServingSim: A HW/SW Co-Simulation Infrastructure for LLM Inference Serving at Scale](https://mlsanity.com/paper/2408.05499/) (2024)
- [Electricity Demand and Grid Impacts of AI Data Centers: Challenges and Prospects](https://mlsanity.com/paper/2509.07218/) (2025)
- [Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces](https://mlsanity.com/paper/2305.14516/) (2023)
- [From Prompts to Power: Measuring the Energy Footprint of LLM Inference](https://mlsanity.com/paper/2511.05597/) (2025)
- [Rethinking Machine Learning Collective Communication as a Multi-Commodity Flow Problem](https://mlsanity.com/paper/2305.13479/) (2023)
- [LLMServingSim 2.0: A Unified Simulator for Heterogeneous and Disaggregated LLM Serving Infrastructure](https://mlsanity.com/paper/2602.23036/) (2026)
- [TACOS: Topology-Aware Collective Algorithm Synthesizer for Distributed Machine Learning](https://mlsanity.com/paper/2304.05301/) (2023)
- [An Efficient, Reliable and Observable Collective Communication Library in Large-scale GPU Training Clusters](https://mlsanity.com/paper/2510.00991/) (2025)
- [Towards a Standardized Representation for Deep Learning Collective Algorithms](https://mlsanity.com/paper/2408.11008/) (2024)
- [Towards Easy and Realistic Network Infrastructure Testing for Large-scale Machine Learning](https://mlsanity.com/paper/2504.20854/) (2025)
- [ForestColl: Throughput-Optimal Collective Communications on Heterogeneous Network Fabrics](https://mlsanity.com/paper/2402.06787/) (2024)
- [ATLAHS: An Application-centric Network Simulator Toolchain for AI, HPC, and Distributed Storage](https://mlsanity.com/paper/2505.08936/) (2025)
- [Maya: Optimizing Deep Learning Training Workloads using GPU Runtime Emulation](https://mlsanity.com/paper/2503.20191/) (2025)
- [Scalable Synthesis of distributed LLM workloads through Symbolic Tensor Graphs](https://mlsanity.com/paper/2511.10480/) (2025)

## Abstract (arXiv)

Distributed machine learning (ML) is a key paradigm for today's large-scale artificial intelligence applications. As model inference arises as an important use case, faithful modeling of latency-sensitive collective communication has never been more important. Capturing the device architecture and modeling control and data paths at high fidelity is therefore a necessity today. Having a common, detailed representation for distributed ML infrastructure is also crucial. We revisit the promising open-source, community-driven simulator: ASTRA-sim. In this work, we identify limitations of the current ASTRA-sim simulator and augment it with new features. To this end, we enable fine-grained, high-fidelity simulation with a standardized infrastructure representation, opening new design space exploration opportunities. We propose the simulation at cache-line-sized load-store granularity, with a detailed graphics processing unit (GPU) execution model, to balance simulation scalability and fidelity. We also introduce InfraGraph, a standardized representation to capture distributed ML network infrastructure in detail. Using the updated ASTRA-sim 3.0 simulator, we showcase interesting design space explorations for designing optimized collective algorithms, network requirements, and GPU architectures.