ASTRA-sim 3.0: Next-Level Distributed Machine Learning Simulations via High-Fidelity GPU and Infrastructure Modeling William Won∗1 , Jinsun Yoo∗2 , Tuan Ta∗1 , Moumita Dey∗1 , Andy Balogh3 , Pradosh Datta3 , Furkan Eris1 , Conor Green1 4 , Winston Liu3 , Changhai Man2 , Kingshuk Mandal3 , Amos Rai3 , Vinay Ramakrishnaiah1 , Ruchi Shah1 , David Sidler1 , Harsh Sikhwal3 , Hanjiang Wu2 , Tushar Krishna†2 , and Bradford M. Beckmann†1
arXiv:2606.10440v1 [cs.DC] 9 Jun 2026
1 AMD Research and Advanced Development
2 Georgia Institute of Technology
Abstract 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.
1
Introduction
Today, the importance of foundation artificial intelligence (AI) models needs no emphasis. Generative AI models, including large language models (LLMs), have been adopted in numerous domains and have been scaling at an unprecedented rate. On the training front, the compute requirements for frontier models have been increasing 4–5× per year on average [12], with LLMs comprising trillions of parameters [10, 13]. Inference outpaces the already massive training requirements, even taking 80–90% of energy usage and cost throughout the model life cycle [3, 36]. To accommodate such scales, the use of distributed machine learning (ML) is necessitated; state-of-the-art data centers contain tens to even hundreds of thousands of graphics processing units (GPUs) [35, 50, 52]. Large-scale distributed ML requires massive resources—not only to train models, but also to service the increasing demands for inference [4, 34]. Considering these needs, optimizing for the performance and cost of distributed ML is crucial. However, the co-design space of distributed ML is vast and complex [49]. It comprises multiple design stacks—AI models, parallelization strategies, GPU * These authors contributed equally to this work. † Equal advising. Correspondence to: William Won <[email protected]>, Tushar Krishna <[email protected]>, Bradford M. Beckmann <[email protected]>.
3 Keysight
4 Purdue University
architectures, and communications between GPUs—and the interplay across them [39]. Navigating this intertwined design space requires a systematic and methodical approach. Therefore, implementing an infrastructure to capture this full-stack design space of distributed ML remains pivotal research. We revisit an existing solution: ASTRA-sim [39, 49]. It is an eventdriven, full-stack distributed ML simulator. We target ASTRA-sim since it is an open-source, community-driven effort and has been adopted by multiple works in the field [6, 7, 31, 53, 54]. ASTRAsim 1.0 first introduced this simulation infrastructure, but supported only a limited set of training workloads and network topologies [39]. ASTRA-sim 2.0, a subsequent update, identified such issues and improved the simulator [49]. Notably, ASTRA-sim 2.0 incorporates graph-based workload representation [43] and extends the scope of network topologies. Despite greater flexibility, considering today’s challenges associated with distributed ML, ASTRA-sim 2.0 still comes with limitations that were overlooked at the time. In this work, we aim to further identify such restrictions and implement new features in the ASTRA-sim framework. To this end, we target capturing the broader design space with improved accuracy. Firstly, we point out the limitations of ASTRA-sim’s collective communication modeling capabilities. For both training and inference jobs, collective communication is the dominant operation [5, 46], yielding multiple fast-paced research efforts. They are spread across new collective algorithm designs [8, 26, 30] and automated collective algorithm synthesis [2, 20, 25, 28, 40, 45, 48, 56], to list a few. However, ASTRA-sim only supports a limited set of predefined, textbook collective algorithms: ring [44], all-pairs (i.e., direct) [39], double binary tree [22], and recursive halvingdoubling [44]. Although collective communication is on the critical path, ASTRA-sim cannot capture the latest research endeavors without manual implementation. Secondly, ASTRA-sim lacks device modeling. It completely overlooks the control path and resource contention. Capturing such effects is especially pivotal today, when inference tasks have become as important, if not more so, as training executions. Inference jobs are latency-sensitive [55], making accurate latency modeling essential. This requires faithful modeling of fine-grained architectural details. For example, assume a GPU wants to write a chunk of data to a remote GPU. This actually consists of multiple steps: (i) a compute unit1 (CU) loads cache-line-sized (e.g., 128 B [1, 33]) 1We use OpenCL/AMD and NVIDIA terminology interchangeably, though the former is
preferred (e.g., compute unit vs. streaming multiprocessor, workgroup vs. threadblock, wavefront vs. warp).
Won et al.
ASTRA-sim 3.0 Chakra Execution Trace (Workload and Parallelization)
Workload Layer Parsing and dispatching kernel-level information from Chakra ET
NoC-Level Network Simulators
System Configuration (e.g., CU specifications)
Custom Collective Algorithm (in MSCCL++)
InfraGraph (Common infrastructure representation)
Parser
System Layer
GPU Model
Translate kernel-level information into simulatable events by the network layer
Simulates fine-grained load/store/compute operations at CU-level granularity
Simple-NoC
Garnet/Gem5
...
Translates kernel-level representation in fine-grained load/store/compute operations
Network Layer Simulate ASTRA-sim events via the network simulations backend (through the Network APIs)
Scale-Out Network Simulators Network APIs
Simple
ns-3
HTSim
...
Figure 1: Overview of the ASTRA-sim 3.0 infrastructure. New and improved components are marked with bold red borders. data from the local high-bandwidth memory (HBM) to its register file (ii) the CU then writes the data to the input/output (I/O) port of the socket (iii) the network transfers the cache-line-sized data to the remote GPU’s I/O port (iv) the remote GPU writes the received data to the destination HBM. These steps are repeated at cache-line granularity until all data chunks are written to the destination. However, the current ASTRA-sim simulates this as a single network transfer, blocking accurate latency estimation. The problem is exacerbated by other important missing control path modeling, including memory fences, intra-GPU barriers, and interGPU semaphores. Multiple GPU kernels fight for the limited CU resources, affecting the latency as well. All such effects must be captured for correct latency modeling. To properly simulate such fine-grained operations, the network model should support network-on-chip (NoC)-level details. However, the current ASTRA-sim infrastructure only assumes coarsegrained, inter-GPU communications. It neglects on-chip transfers such as a CU loading cache-line-sized data from the local HBM channel. The network simulation backend should also be upgraded to support the modeling of the architectural details of a GPU socket. Finally, the community requires a common, detailed representation of distributed ML network infrastructure. Currently, different network simulation backends all require unique input formats. As different simulators capture infrastructure at distinct levels with different details, this fragmentation not only hinders accurate infrastructure modeling but also prohibits the community from exchanging their infrastructure details. In this work, we introduce ASTRA-sim 3.0: taking distributed ML simulations to the next level via high-fidelity modeling. Figure 1 summarizes the improved ASTRA-sim 3.0 simulation infrastructure. ASTRA-sim 3.0 implements new features to eliminate the limitations and to faithfully model today’s distributed ML systems. ASTRA-sim 3.0 supports arbitrary customized collective algorithms through MSCCL++ representations [41]. ASTRA-sim 3.0 purposely represents workloads at Load-Store granularity, capturing finegrained details while remaining scalable. A new GPU Model simulates them at cache-line-sized granularity with CU, workgroup, and wavefront-level modeling, just as real GPUs do. Finally, we build InfraGraph, a graph-based backend-agnostic representation of the physical system. Users can exchange and reuse a single InfraGraph description across network backends in ASTRA-sim 3.0.
2 Background 2.1 ASTRA-sim ASTRA-sim [39, 49] is an open-source distributed ML simulator. The high-level architecture is summarized in Figure 1. It consists of three building-block layers: workload, system, and network. The workload layer captures the design space of the target model and its parallelization, whereas the network layer simulates the network transfers. The system layer intermediates between the two layers, for example, estimating compute kernel runtime or breaking down a collective communication kernel into chunk-granularity transfers using predefined collective algorithms. ASTRA-sim 1.0 is the introductory version but lacked flexibility in all three layers [39]. It only supported a predefined set of training workloads, collective algorithms, and network topologies (specifically, two-dimensional switch and torus) through Garnet, a network model originally developed in gem5 [29]. Later, an application programming interface (API) to plug in arbitrary network simulators was introduced, and ns-3 [18] was incorporated as another simulation backend [38]. ASTRA-sim 2.0 addressed some flexibility limitations [49]. For the workload layer, ASTRA-sim 2.0 leverages MLCommons Chakra execution trace (ET) [43] to support arbitrary workloads. For the network layer, ASTRA-sim 2.0 introduces a new 𝛼–𝛽-model-based [19] Simple network simulator.2
2.2
Graphics Processing Unit
We briefly introduce the background on GPUs, the dominant compute devices in distributed ML [16, 23]. 2.2.1 Programming Model. A kernel is a function to be executed on a GPU, and comprises many threads. Each thread executes the same instruction but with different data, making the GPU follow the single-instruction, multiple-data (SIMD) paradigm. Threads in a kernel are organized into workgroups (i.e., threadblocks), a composition of threads (e.g., 1,024 threads). All threads in a workgroup are mapped and executed on one CU (i.e., streaming multiprocessor (SM)). Within a workgroup, threads are further grouped into a smaller unit called wavefronts (i.e., warps), a lock-step execution unit within a CU (e.g., 32 threads). Figure 2 depicts a GPU kernel with four workgroups, each containing two wavefronts. 2We name it Simple in this work (previously called Analytical), following its resem-
blance to the Simple network simulator in gem5 [14].
ASTRA-sim 3.0: Next-Level Distributed Machine Learning Simulations via High-Fidelity GPU and Infrastructure Modeling
Kernel 0 Workgroup 1 Wavefront 1
Wavefront 2
Workgroup 2 Wavefront 1
Workgroup 3
Wavefront 2
CU
CU
CU
CU
Wavefront 1
Wavefront 2
CU
CU
CU
CU
CU
CU
{ ... "gpus": [ { "id": 0, "threadblocks": [ { id: 0, "ops": [ "put", "signal", "wait" ]}, { id: 1, "ops": [ "signal", "put", "wait" ]} ] }, { "id": 1, "threadblocks": [ { id: 0, "ops": [ "put", "nop", "get" ]}, { id: 1, "ops": [ "nop", "signal", "wait" ]} ] } ] ...}
Workgroup 4 Wavefront 1
Wavefront 2
GPU
Figure 2: Abstract view of GPU program and architecture.
Figure 3: A simplified example of MSCCL++ JSON collective algorithm representation.
3 2.2.2 GPU Architecture. GPU hardware consists of many CUs to execute multiple workgroups in parallel. When a kernel is dispatched to the GPU, each workgroup is mapped onto a CU. The CU then executes threads in the wavefront in lock-step. Figure 2 shows that the four workgroups in the kernel are being executed in parallel by four CUs.
2.3
Collective Communication
The model and/or data are dispersed across many devices in distributed ML, requiring that the devices frequently synchronize to execute the full workload. Such communication takes place in the form of collective communication patterns defined in the message passing interface (MPI) [32]. A collective algorithm defines how individual data chunks are transferred over a network topology to execute the collective patterns. Textbook collective algorithms are well-defined, basic collective algorithms that are pervasively used in today’s collective communication libraries (CCLs). Ring [44], all-pairs (i.e., direct) [39], double binary tree [22], and recursive halving-doubling [44] are examples of predefined textbook collective algorithms used by CCLs today.
2.4
MSCCL++ Custom Collective Representation
As collective communication becomes a crucial part of distributed ML, many research efforts have been put into the domain; among them is the representation of collective algorithms. Notably, MSCCLang introduces a domain-specific language (DSL) to represent arbitrary collective communication algorithms [9]. Customized collective algorithms can be written in a Python-based DSL. Later, MSCCL++ was introduced to generalize the DSL and to represent a more flexible set of collective algorithms [41]. For instance, MSCCL++ supports the use of one-sided put/get operations, allows one workgroup to talk to multiple remote GPUs simultaneously, and captures control dependencies via barriers and semaphores. MSCCL++ compiles the DSL into a custom JavaScript object notation (JSON) file, capturing workgroup-level operations of individual GPUs. Figure 3 captures a simplified view of the JSON file that represents a collective operation of two GPUs. Each GPU comprises two workgroups, each with three GPU operations.
Motivation
Here, we identify the limitations of the current ASTRA-sim. We discuss why having such modeling capabilities is important for solving today’s distributed ML challenges.
3.1
Customized Collective Algorithms
The current ASTRA-sim infrastructure completely lacks customized collective algorithm modeling capabilities. As Section 1 discusses, collective communication has become the major bottleneck in distributed ML. For training tasks, collective sizes are often large, and the target is to optimize for throughput, while inference communication is smaller in size and latency-sensitive. Such different natures yield distinct approaches and techniques for collective optimization. To list a few, even for textbook collective algorithms, experts fine-tune the number of workgroups, the size of individual chunks, communication protocols, or communication primitives (e.g., twosided vs. one-sided transfers). Collective algorithm synthesizers automatically generate topology-aware collective algorithms. As the system scales, fault-tolerant collective algorithm design also becomes an important research angle. Ultimately, the existence and adoption of MSCCL++ representation itself highlight the necessity to execute customized collective algorithms. However, ASTRA-sim infrastructure only supports a limited set of predefined textbook collective algorithms, with very limited design parameters (e.g., number of chunks per kernel). ASTRA-sim cannot simulate today’s fast-changing collectives without extensive manual implementation.
3.2
Fine-Grained GPU Modeling
The existing ASTRA-sim infrastructure does not capture fine-grained device modeling. With the massive adoption of AI workloads, the community is shifting gears from training to ML inference. As inference workloads are highly latency-sensitive, correct latency modeling is essential. Faithful latency modeling necessitates high-fidelity modeling of the data and control paths of the device. As articulated in Section 1, even a simple network put operation actually involves multiple repetitive steps working at cache-line granularity: (i) after making sure the destination buffer is ready, (ii) a CU reads cache-line-sized data to the register file, and (iii) initiates the network transfer of the read value; then (iv) the remote GPU stores the received data to the destination. However, none of these fine-grained operations are modeled in ASTRA-sim. This is not to mention that multiple workgroups from different communication and compute kernels contend
Won et al.
0.5 us
5 us
200
256 GiB/s
100 0 1000
Protocol
Workgroup
Workgroup 1
Workgroup 2
Wavefront 1
Wavefront 2
Workgroup 3
Workgroup 4
Wavefront 3
Wavefront 4
LL
Wavefront
GPU Operation
NopOp
Load
MemCopyOp
WaitCnt
BarrierOp
Store
Simple
750
1 TiB/s
Transfer Bandwidth (GiB/s)
Kernel
500 250 512 MiB
8 MiB
64 MiB
1 MiB
128 KiB
2 KiB
16 KiB
256 B
512 MiB
8 MiB
64 MiB
1 MiB
128 KiB
2 KiB
16 KiB
256 B
0
Figure 5: Abstract view of a GPU kernel broken down into fine-grained Load-Store granularity in ASTRA-sim 3.0.
Transfer Size
Figure 4: Analytical analysis of LL versus Simple transfer bandwidth, using different link latency and bandwidth. for the limited CU and memory resources, further impacting the overall latency. Figure 4 is a quantitative motivation for why such accurate latency modeling is required. Most CCLs in use today provide two communication protocols: low-latency (LL) and Simple. The Simple protocol can effectively leverage 100% of network bandwidth but comes with the tradeoff of synchronization before and after the transfer. The LL protocol eliminates the synchronization for latency optimization with the tradeoff of 50% link bandwidth. Intuitively, LL is preferred for small transfers, and Simple starts to outperform as the transfer size increases. Using the analytical 𝛼–𝛽 modeling with arbitrary latency and bandwidth values, we plotted the transfer bandwidth of the two protocols with different transfer sizes. When we underestimate the transfer latency (as 0.5 𝜇s), the Simple protocol starts to outperform LL much faster. For 256 GiB/s link bandwidth, the Simple protocol outperformed LL at a 512 KiB transfer size. However, if we overestimate the link latency (as 5 𝜇s), the Simple protocol performance saturates at much larger transfer sizes; it only outperformed the LL protocol at a 2 MiB transfer size. As the link bandwidth increases to 1 TiB/s, the gap is even larger: 4 MiB versus 16 MiB. This simple analytical analysis underscores that improper latency modeling can lead to wrong design conclusions, which becomes especially crucial as the network becomes more performant. High-fidelity latency modeling through fine-grained GPU modeling is therefore necessitated in the simulation infrastructure.
3.3
Common Infrastructure Representation
To accommodate different use cases (e.g., fidelity versus scalability), ASTRA-sim extends to multiple network simulation backends via common network APIs. However, users should define the physical system topology in a format specific to the target network backend. Examples include routing configurations, the duration of a message, or whether to model congestion at a given point. Through our experience using ASTRA-sim, we found the lack of a global representation format to be problematic. Firstly, any physical system description is tied to a specific network backend. This fragmented approach forces users to manually rewrite infrastructure descriptions across network backends, not to mention the time to understand the semantics of the format, which is not intuitive and requires per-backend knowledge. Secondly, it hinders the community from sharing and exchanging the exact
infrastructure details. Existing research often describes physical systems with high-level verbiage such as “two-tiered fat-tree topology" or through simple figures. However, this may not necessarily capture the full detail of the topology, making it vague when translating into a description that ASTRA-sim backends can understand. If the community had a common, standardized means to represent their physical topology information in a neutral format, it could facilitate the community in sharing infrastructure details and reproducing the simulation results. We envision a potential beyond ASTRA-sim in standardizing the infrastructure representation. In short, the above motivates a backend-agnostic format that can represent arbitrary physical system topologies. We motivate the need for a single format that, once defined, not only captures the fine-grained details of the network infrastructure, but also can be easily shared across the research community. Such representation can be automatically translated to the backend-specific format for ASTRA-sim.
4
ASTRA-sim 3.0
Here, we explain the new features added to ASTRA-sim 3.0. We also articulate their implementation in detail.
4.1
Fine-Grained Workload Representation
We propose to model the GPU device at the Load-Store granularity. The critical operations in distributed ML are the control and data paths over the on-chip and inter-GPU networks. They can effectively be modeled by capturing the fine-grained Load-Store instructions. Simulating the exact GPU binary at the finest granularity contains unnecessarily detailed information and hurts simulation scalability. A high-level overview of the GPU kernel decomposed in ASTRAsim 3.0 is shown in Figure 5. It captures the execution in the manner in which the GPU programming model does: a workgroup comprises multiple wavefronts, each running multiple GPU operations in sequence, where each GPU operation itself is ultimately a sequence of Load-Store primitives. Below, we articulate the finegrained workload representation in ASTRA-sim 3.0 in a bottom-up approach. 4.1.1 GPU Instructions. As we proposed, we define primitive GPU instructions at the Load-Store granularity. These GPU instructions become the unit of simulation in ASTRA-sim 3.0. The GPU instructions can be grouped into three: (i) data instructions, (ii) control instructions, and (iii) others. For these Load-Store, the source or destination memory location may reside either locally or remotely.
ASTRA-sim 3.0: Next-Level Distributed Machine Learning Simulations via High-Fidelity GPU and Infrastructure Modeling
Data Instructions. They represent the movement of data between the CU and memory, as defined below: • Load: A CU loads a chunk of data from memory to the register file. • Store: A CU stores a chunk of data from the register file to memory.
may itself be a meaningful operation. An example is an operation to halt execution until a certain condition is met. We define two such operations for GPU execution control without any memory instructions. • ReduceOp: A wrapper of the Reduce instruction. Represents all arithmetic and logical operations without memory operations. • NopOp: This operation halts the execution of a workgroup until all its wavefronts reach this operation (i.e., __syncthread operation). • BarrierOp: While NopOp enforces intra-workgroup synchronization of wavefronts, BarrierOp enforces inter-workgroup synchronization.
Control Instructions. These instructions also create load and store requests. However, we differentiate these from the data operations to emphasize their nature as control path instructions. They are listed below: • SemaphoreAcquire: A CU loads a semaphore value and checks if it is released by another CU. • SemaphoreRelease: A CU stores a semaphore value to mark that it is released and can be acquired by another CU. Other Instructions. Finally, we add two more primitive instructions, as below, to capture non-memory executions. • Reduce: Abstracts all other arithmetic or logical operations taking place in the CU. • Waitcnt: Puts the CU on wait until the number of in-flight Load and Store drops to a specific threshold, to control intra-wavefront memory dependencies. 4.1.2 GPU Operations. While these primitive Load-Store instructions capture the fine-grained behaviors of a GPU, they are too fine-grained to capture the logical operation the programmer leverages. To mitigate this gap, we introduce GPU operations. A GPU operation is a sequence of primitive GPU instructions, and it denotes a meaningful, functional programming unit. Examples include loading data from a range of memory addresses or synchronizing all threads within a workgroup. For the scope of this work, ASTRA-sim 3.0 implements several data- and control-related GPU operations. However, any new GPU operation can be easily implemented, as it is simply a sequence of multiple Load-Store GPU instructions. Data Operations. These operations are used to load and store a range of data. We predefine three data operations as below:
4.1.3 Workgroup and Wavefront. Now that we have defined the logical execution operations a GPU can execute, defining a workgroup and wavefront is straightforward. A workgroup is simply a sequence of GPU operations, executed over a single CU. A workgroup comprises multiple wavefronts, a unit of lock-step execution in the CU. For data operations, all wavefronts execute the Load and Store instructions, simulating each wavefront processing different ranges of memory addresses. For control path operations, on the other hand, we implement only wavefront zero to issue the Load and Store instructions, assuming a control message (i.e., reading a semaphore value) is single cache-line-sized. 4.1.4 Kernel. Finally, a kernel is a set of workgroups. When a kernel is dispatched to a GPU, each workgroup is mapped to an individual CU and executes in parallel. To summarize, as in Figure 5, (i) a GPU instruction is a primitive Load-Store instruction, (ii) a GPU operation is a sequence of GPU instructions, denoting a logically meaningful execution unit, (iii) a workgroup comprises multiple GPU operations in sequence, organized in multiple wavefronts, and (iv) a kernel is a group of workgroups that run in parallel on multiple CUs.
4.2
MSCCL++ Custom Collective Support
Now that ASTRA-sim 3.0 represents the workload in a fine-grained representation, this easily enables the simulation of custom collective algorithms leveraging the MSCCL++ representation. As explained in Section 2.4, MSCCL++ JSON representation captures arbitrary collective algorithms in per-GPU, per-workgroup operations such as put or copy. Therefore, we simply implement a straightforward translator to represent MSCCL++ operations in ASTRAsim 3.0 GPU operations. For example, a MSCCL++ put operation (writing chunks of data from a local GPU to a remote GPU) is translated into MemcpyOp. The same applies to the get and copy operaControl Operations. These two control operations wrap the semaphore tions. MSCCL++ inter-GPU control operations, signal and wait, instructions, as shown below: are parsed into SemaphoreReleaseOp and SemaphoreAcquireOp operations, respectively. • SemaphoreAcquireOp: Issues a SemaphoreAcquire instruction to acquire the semaphore. • SemaphoreReleaseOp: Issues a SemaphoreRelease instruc4.3 End-to-End Workload Execution tion to release the semaphore. ASTRA-sim 3.0 inherits the end-to-end workload simulation from • LoadOp: A wrapper of the Load instruction. Loads a range of data from memory to the CU. • StoreOp: A wrapper of the Store instruction. Stores a range of data from the CU to memory. • MemcpyOp: Represents memory-to-memory copy operations. Data is first loaded to the CU by the Load instruction, followed by the Waitcnt instruction to enforce a memory fence. Finally, the Store instruction writes the loaded data from the CU to the destination memory space.
Other Operations. We define ReduceOp to wrap the Reduce instruction to capture all non-memory arithmetic and logical operations. Interestingly, we note that dispatching no GPU instructions
ASTRA-sim 2.0 via MLCommons Chakra ET adoption. As shown in Figure 6, Chakra ET captures workloads at kernel granularity: a trace is a composition of compute and communication kernels and
Won et al.
Compute Kernel
Compute Parser Compute Kernel
Communication Kernel
Communication Kernel
Communication Kernel
(a) Chakra ET representation
Workgroup 2
LoadOp
LoadOp
ReduceOp
ReduceOp
Workgroup 1
Workgroup 2
MemcpyOp
MemcpyOp
Semaphore
MemcpyOp
Load
Load
Load
Load
WaitCnt
Waitcnt
Store
Waitcnt
Store
Store
Waitcnt
Store
Waitcnt
Store
Load
Figure 7: A MemcpyOp operation unrolled by four times.
(b) Parsers translate each kernel into fine-grained workload representation
Figure 6: End-to-end workload simulation flow through ASTRA-sim 3.0.
the dependencies between them. ASTRA-sim 2.0 dispatches coarsegrained events through parsing Chakra ET. Instead, ASTRA-sim 3.0 implements parsers to decompose such kernels into fine-grained workload representations. For communication kernels, we leverage the MSCCL++ parser discussed in Section 4.2. Likewise, compute kernels can be decomposed into multiple workgroups, each executing ReduceOp operations. ASTRA-sim 2.0 used very basic hardware models to simulate resource contention, such as dispatching one kernel at a time. However, in ASTRA-sim 3.0, all kernels are decomposed into the common, fine-grained representation and simulated by a single model. Therefore, resource contention between multiple compute or communication kernels is naturally captured without such arbitrary restrictions.
4.4
GPU Operation
Unroll by 4
Communication Kernel
MSCCL++ Parser Compute Kernel
Workgroup 1
GPU Model
So far, we have discussed how ASTRA-sim 3.0 represents the target workload in a fine-grained, Load-Store instruction representation. ASTRA-sim 3.0 implements a new execution model to simulate this fine-grained representation in detail to capture the actual program execution. Below, we explain this execution model in a top-down approach. 4.4.1 GPU Model. As the name suggests, the GPU Model abstracts a physical GPU and consists of multiple CUs. The ASTRA-sim 3.0 simulation starts by dispatching a kernel (parsed in the fine-grained representation) to GPU Model. The GPU Model iterates over the workgroups and maps each workgroup to a (free) CU in a roundrobin order, thereby modeling CU resource conflicts. The number of workgroups that each CU can accommodate can be tuned by the user. 4.4.2 Compute Unit. When a workgroup is dispatched to a CU, it executes the wavefronts from the workgroup. The CU chooses a wavefront that is ready (i.e., not stalled) and executes the GPU operations in sequence. During the execution, if a wavefront stalls due to control path operations, the CU executes another ready wavefront, modeling wavefront-level parallelism. If all GPU operations of a wavefront are finished, then the wavefront is marked as done. When all wavefronts of a workgroup are completed, the workgroup is marked as finished and retires. The items below explain how the non-memory-based GPU operations execute in the CU:
• NopOp: As mentioned in Section 4.1.2, this operation does not have any GPU instructions. Instead, when executing this operation, the CU simply puts the wavefront in a stalled state, and checks if all other wavefronts in the workgroup are also stalled. If so, NopOp completes and all wavefronts are marked as ready to execute the next operation. • BarrierOp: The execution follows the same logic as NopOp, but enforces inter-workgroup synchronization rather than intra-workgroup synchronization. • ReduceOp: Occupies the CU for a certain number of simulation cycles to mimic the arithmetic/logical operations, then retires the operation. Other GPU operations comprise Load or Store instructions. For example, LoadOp requires a Load instruction to load a value from memory. Likewise, SemaphoreAcquireOp wraps a SemaphoreAcquire instruction, which dispatches a Load instruction to check the semaphore value. When executing these operations, the CU injects cache-linesized network requests, namely Wavefront Requests, into the network simulator in each cycle. 4.4.3 Wavefront Request. A Wavefront Request is a single cacheline-sized (e.g., 128 B) transfer that a CU dispatches to the network. In ASTRA-sim 3.0, the Wavefront Request is the unit transfer that the network backend simulates. The network layer receives these Wavefront Requests and simulates the local (on-chip) and remote (scale-up and scale-out) network traffic. When a CU executes the GPU operation of a wavefront, it generates multiple Wavefront Requests depending on the operation. • SemaphoreAcquireOp and SemaphoreReleaseOp: Acquiring and releasing a semaphore requires reading and writing the semaphore value, respectively. As we assume each semaphore fits in a single cache line, executing these operations generates one Wavefront Request from wavefront zero. • LoadOp, StoreOp, and MemcpyOp: They generate multiple Wavefront Requests to read or write a memory span. 4.4.4 Loop Unrolling. When a CU executes a wavefront, it can dispatch one cache-line-sized Wavefront Request in each cycle. A GPU may facilitate intra-wavefront, instruction-level parallelization by dispatching multiple in-flight memory requests simultaneously. In ASTRA-sim 3.0, we implement this via tunable loop unrolling. Figure 7 exemplifies the unrolling of a MemcpyOp operation. MemcpyOp requires the repetition of (i) loading cache-line-sized data, (ii) making sure the data has been loaded by the Waitcnt instruction, then (iii) storing the data to the destination. Instead of repeating the three primitive instructions for each cache line, Figure 7 unrolls it
ASTRA-sim 3.0: Next-Level Distributed Machine Learning Simulations via High-Fidelity GPU and Infrastructure Modeling
GPU 1
GPU 2
HBM
GPU 1
GPU 2
GPU 3
GPU 4
HBM
HBM
HBM
HBM
HBM
HBM
HBM
CU
CU
I/O
CU
CU
CU
CU
I/O
CU
CU
I/O
I/O
I/O
I/O
CU
CU
I/O
CU
CU
CU
CU
I/O
CU
CU
HBM
HBM
HBM
HBM
HBM
GPU 3
(a) Simple network: each endpoint models a GPU
HBM
HBM
HBM GPU 4
(b) Updated Simple network: captures and simulates detailed NoC-level transfers
Figure 8: Updated Simple network simulation backend. by a factor of four. The unrolling enables overlapping four Load Wavefront Requests, capturing intra-wavefront parallelism.
4.5
Detailed Simple Network Modeling
With the new custom collective and GPU modeling capabilities, ASTRA-sim 3.0 is equipped with detailed, fine-grained simulation opportunities with NoC-level requests and cache-line-sized transfers. To better accommodate this opportunity, we upgrade the Simple network simulation backend from ASTRA-sim 2.0 to simulate NoC-level details alongside the inter-GPU network. Note that ASTRA-sim 3.0 still supports other network simulation backends with detailed NoC models through the network API. The high-level insight is depicted in Figure 8. The old Simple network model, shown in Figure 8(a), used each GPU as the building block. The network representation was restricted to simple GPU connections and simulated through coarse-grained transfers. The updated Simple network is shown in Figure 8(b). Now the target network topology is represented with NoC-level details, such as CUs, HBM channels, and I/O ports. As the GPU Model of the ASTRA-sim 3.0 infrastructure dispatches fine-grained Wavefront Request, the updated Simple with the NoC-level representation is able to simulate both local and remote operations. Memory channels are added to Simple to model the NoC and memory traffic. The network representation of Simple uses a hierarchical file-based representation and modular routing function to allow scalability and quick interchanging of network topologies or routing.
4.6
Figure 9: Clos fabric generated and visualized using InfraGraph blueprint and visualizer.
Infrastructure as a Graph
We present InfraGraph, a standard, portable representation for AI and high-performance computing (HPC) network infrastructure. InfraGraph formalizes infrastructure topology as a directed, attributed graph in which vertices represent hardware like GPUs, network interface controllers (NICs), and storage, while edges represent the connections between the hardware components. Property annotations such as bandwidth enhance the expressiveness of the graph. This representation is based on the principle that complex infrastructure topologies can be naturally expressed through graph constructs. Users define a graph by describing reusable infrastructure objects and topology relationships and programmatically expanding them into a complete graph representation. Reusing components allows
a compact description to expressively represent a fully expanded graph, while programmatic graph construction enables automatic workflows. 4.6.1 Device Description. We first define the primitives used to describe a Device: • Component: Hardware unit within a device such as a CPU, GPU, or peripheral component interconnect express (PCIe) bridge. • Link: Named connection container with physical properties such as bandwidth or latency. 4.6.2 InfraGraph Graph Construction. A core design principle of InfraGraph is to make the definition compact by reusing modular definitions. Rather than manually defining every endpoint and interconnection in a large-scale AI cluster, users describe reusable infrastructure objects that can be programmatically expanded into a complete graph. For example, in a scale-out topology across multiple host servers, a single host can be represented as a Device where hardware Component such as a CPU, GPU, or PCIe bridge are represented as vertices and PCIe Edges are the edges. Instead of repeatedly listing the Components and Edges multiple times for all hosts, a user can elect to programmatically construct them. The below primitives enable large construction: • Device: Subgraph template for device hardware, containing Component and Edges. • Instance: Device instantiation alias. • Infrastructure: Top level graph container. • Device.Edge: Edge of an infrastructure graph defined by two Device endpoints and a connecting Link. 4.6.3 Blueprints. We also offer pre-built, composable templates for common hardware platforms and network fabrics. Device blueprints define the internal hardware structure of a single platform—components, links, and intra-device edges—and serve as reusable building blocks. Fabric blueprints compose device instances into full network topologies. These blueprints are parameterized (e.g., the number of hosts) and automatically instantiated. For example, SingleTierFabric fabric blueprint implements a flat single-switch-layer topology for small-scale deployments, while ClosFatTreeFabric produces a scalable hierarchical topology parameterized by switch port count and network depth, automatically computing switch counts and wiring all links per the standard CLOS construction.
Algorithm
400 get−based 300
put−based
200 100 0
256 MiB
128 MiB
64 MiB
32 MiB
8 MiB
16 MiB
4 MiB
2 MiB
1 MiB
0
512 KiB
100
With Arbitration
16 KiB 32 KiB 64 KiB 128 KiB 256 KiB 512 KiB 1 MiB 2 MiB 4 MiB 8 MiB 16 MiB 32 MiB 64 MiB 128 MiB 256 MiB
put−based
200
Without Arbitration
16 KiB 32 KiB 64 KiB 128 KiB 256 KiB 512 KiB 1 MiB 2 MiB 4 MiB 8 MiB 16 MiB 32 MiB 64 MiB 128 MiB 256 MiB
get−based
300
Collective Bandwidth (GiB/s)
Algorithm
400
256 KiB
Collective Bandwidth (GiB/s)
Won et al.
Output Buffer Size
Input Buffer Size
Figure 10: Simulated collective performance of get- and putbased Reduce-Scatter with 32 GPUs.
Figure 11: Simulated collective bandwidth of get- and putbased All-Gather with 16 GPUs, with and without arbitration between control and data messages.
5 4.7
InfraGraph Toolchain
4.7.1 Translator. We build a translator that takes the InfraGraph description and automatically translates this into a physical system description that each network backend in ASTRA-sim can understand. The translator has three backend translators for the publicly available network backends of ASTRA-sim: Simple, ns-3, and HTSim. Each backend translator automatically derives its required parameters from the annotated graph. The HTSim translator infers fat-tree structural parameters from the topology; ns-3 assigns identifiers to GPU ranks, NICs, and switches and reads per-link properties from edge annotations; and the Simple translator additionally detects topology patterns to decompose large node counts into multi-dimensional groups for collective communication modeling. In all cases, the same InfraGraph description produces valid configurations across all backends, enabling direct cross-backend comparison under identical infrastructure assumptions. 4.7.2 Visualizer. The visualizer receives an InfraGraph and automatically generates network connectivity plots. This allows users to see whether the network graph they defined is what they indeed intended. 4.7.3 Fully Qualified Graph and InfraGraph Service. InfraGraph expands a topology description into a fully qualified graph composed of infrastructure nodes and edges. Each device, component, port, and link is represented using unique hierarchical identifiers, enabling accurate topology representation, communication path discovery, graph traversal, connectivity analysis, and topologyaware simulation across large-scale AI and HPC infrastructures. An infrastructure node represents the fundamental entity within the infrastructure graph and corresponds to a specific endpoint in the system topology. A node follows the naming convention: <device-instance>.<index>.<component>.<index>. This hierarchical naming structure enables precise identification of components within large-scale heterogeneous systems. An infrastructure edge represents a communication relationship between two infrastructure nodes. Edges define connectivity within the graph and model communication links between components. An example of an edge representation for a switch is: (switch.0.asic.0, switch.0.port.0, pcie) where it is represented as a source node, a destination node, and the link connecting both nodes.
Case Studies
In this section, we run various case studies to showcase how ASTRAsim 3.0 enables new design space exploration opportunities that previous simulators could not capture. We wish to highlight that this section aims to demonstrate various use cases and the applicability of the ASTRA-sim 3.0 simulation infrastructure, not necessarily showcasing quantitative optimization measures of specific products or infrastructures. Therefore, we evaluate a generic GPU architecture with power-of-two configuration values.
5.1
Target GPU Architecture
For the fine-grained GPU modeling case studies, we implement a generic GPU architecture. We model a two-dimensional mesh NoC with 32 routers (8×4) with 1 TiB/s on-chip links, each equipped with four CUs (i.e., 128 CUs total per GPU). The topmost and bottommost NoC routers are also connected to 16 memory channels each, providing 4 TiB/s cumulative memory bandwidth. Likewise, the leftmost and rightmost NoC routers are equipped with four I/O ports each, driving 1 TiB/s cumulative scale-up bandwidth with 1 𝜇s link latency per GPU. In total, each GPU is modeled as a collection of 448 endpoints.
5.2
Collective Algorithm Design
We first demonstrate how ASTRA-sim 3.0 enables a new realm of modeling collective algorithm designs and assists network designers in optimizing the target system, through the fine-grained, faithful control path modeling. Unlike existing infrastructures that don’t support custom collective algorithms, ASTRA-sim 3.0 enables the comparison of using get versus put for data transfers, which imply different synchronization requirements. We evaluate the bandwidth of get- and put-based Reduce-Scatter algorithms over a 32-GPU cluster with 32 workgroups per GPU. Simulated collective bandwidths (i.e., buffer size divided by collective time) are shown in Figure 10. The results show that get-based Reduce-Scatter outperforms put-based Reduce-Scatter for large collectives. Modeling Insight. For put operations, the sender must notify the receiver after the completion of the transfer (i.e., SemaphoreReleaseOp should follow after a remote Store). This incurs unnecessary control traffic and hinders the overlap of data transfer and reduction computation, as the receiver must run SemaphoreAcquireOp before ReduceOp. However, get operations eliminate synchronization. Once the initiator receives the data via a remote Load, that
Output Buffer Size
5.3
GPU Architecture Exploration
ASTRA-sim 3.0 fine-grained modeling is also applicable to GPU architecture design and enables new design space exploration capabilities. For example, ASTRA-sim 3.0 can be used to model and determine the degree of required intra-wavefront, instruction-level parallelism to optimize collective performance. Figure 12 measures the All-to-All performance using 16 GPUs, 60 workgroups per GPU. Modeling Insight. The result indicates that increased instructionlevel parallelism, which enables each CU to dispatch multiple concurrent Wavefront Request, is beneficial for improving bandwidthbound collectives. Notably, (i) the benefit saturates after hitting the maximum number of Wavefront Request each CU can dispatch at a time, and (ii) instruction-level parallelism is not a relevant optimization for latency-bound, small-sized collectives. We also evaluate how the maximum number of Wavefront Request each CU can dispatch impacts the collective bandwidth. Figure 13 summarizes the All-Gather performance with 32 GPUs, 62 workgroups per CU. Modeling Insight. This configuration is a proxy for the register file size, as it is the determining factor of how many cache-line-sized requests a CU can service concurrently. Similar to the loop unrolling observation, (i) register file size does not have a meaningful impact
256 MiB
64 MiB
128 MiB
32 MiB
8 MiB
16 MiB
4 MiB
2 MiB
Figure 13: Simulated All-Gather bandwidth with 32 GPUs, with different numbers of maximum outstanding Wavefront Request limits per CU. 80000
GPU Count
60000
2
40000
4 8
20000
16
0 256 MiB
128 MiB
32 MiB
32 1 MiB 16 MiB
GPU can immediately start the reduction. This enables computecommunication overlap at cache-line granularity. We also compare get- versus put-based collective algorithms for All-Gather, using 16 GPUs with 60 workgroups each. Figure 11 captures the result. Notably, the get-based collective is less performant than the put-based one. Modeling Insight. Unlike Reduce-Scatter, All-Gather does not have a reduction operation, eliminating the benefits of get operations. Rather, as the buffer size becomes large, the control messages (i.e., requests to the remote GPU to send data) are being blocked by data responses. put operations, on the other hand, push the data first to the network, and the responses (i.e., acknowledgment messages after remote Store) have a marginal effect on performance even if blocked. By using fair arbitration of control and data messages, we minimized the performance gap. Different collective algorithms, even as small as using get versus put, have a significant impact and implications for performance and network design requirements and configurations. They exemplify the necessity of having fine-grained control path modeling that ASTRA-sim 3.0 enables.
2048
Output Buffer Size
Simulation Time (s)
Figure 12: Simulated All-to-All performance of varying loop unrolling factors with 16 GPUs.
1 MiB
1024 512 KiB
512
0 256 KiB
256 MiB
64 MiB
128 MiB
32 MiB
8 MiB
16 MiB
4 MiB
2 MiB
1 MiB
512 KiB
256 KiB
64 KiB
128 KiB
32 KiB
32
256
100
64 MiB
16
0
128
200
64 KiB
8
100
300
128 KiB
4
Outstanding Wavefront Requests per CU
32 KiB
200
400
16 KiB
300
2
Collective Bandwidth (GiB/s)
Loop Unrolling Factor
400
16 KiB
Collective Bandwidth (GiB/s)
ASTRA-sim 3.0: Next-Level Distributed Machine Learning Simulations via High-Fidelity GPU and Infrastructure Modeling
64 128
Output Buffer Size
Figure 14: Wall-clock simulation time of All-Gather for 1– 256 MiB output buffer sizes, when modeling 2–128 GPUs. on control-path-dominated, latency-bound collectives, and (ii) the benefit saturates quickly after hitting a specific register file size. These examples highlight the effectiveness of having fine-grained, Load-Store-level architecture modeling that enables new architectural explorations.
5.4
Fine-Grained Simulation Scalability
We also analyze the scalability of the fine-grained, cache-line-sized Load-Store modeling with the updated Simple network’s NoClevel simulation. Specifically, we measure the wall-clock simulation time of All-Gather collective communication with 32 workgroups per GPU, for 1–256 MiB output buffer sizes. We scale the target system to 2–128 GPUs. Notably, a 128-GPU cluster simulates 57,344 (128 GPUs×448) endpoints simultaneously at 128 B granularity. The results are summarized in Figure 14. We also plot the simulation throughput (i.e., how many nanoseconds ASTRA-sim 3.0 can simulate per wall-clock second) in Figure 15. Simulation Insight. The larger the collective size, the larger the number of 128 B-sized Wavefront Requests the simulation requires to be simulated. Therefore, for all scenarios, the simulation time is linear in the output buffer size. Regardless of the number of Wavefront Requests dispatched to the network simulator, the simulator retained similar simulation throughput. Rather, it is determined by the target system scale being modeled.
5.5
Scale-Out Infrastructure Simulation
Finally, we validate ASTRA-sim 3.0’s ability to capture distinct networks via a standardized InfraGraph. We simulate a 1 MB ring AllReduce across eight GPUs interconnected via a Clos fabric shown
800
GPU Count
600
2 4
400
8 16
200
32 256 MiB
128 MiB
64 MiB
32 MiB
16 MiB
8 MiB
0 4 MiB
Simulation Throughput (simulated ns per second)
Won et al.
64 128
Output Buffer Size
Figure 15: Simulation throughput (i.e., simulated nanoseconds per wall-clock second) for 2–128 GPUs.
in Figure 9 using the ns-3 packet-level backend. The system is instantiated from the InfraGraph device blueprint, with per-device and per-link properties bound via the annotation layer. Results are summarized in Table 1. The simulation properly modeled the Clos scale-out network, yielding a minimum flow completion time (FCT) of 11,250 ns and a standalone FCT of 11,857 ns. The infrastructure does not incur any packet drops, confirming that the ring traffic pattern remains lossless within the simulated fabric.
6
Related Work
We largely group previous simulators into two domains: (i) GPU binary simulators and (ii) distributed ML simulators. GPU Binary Simulators. GPGPU-sim [17] and MacSim [15] are cycle-accurate simulators of GPU executables, modeling the complete instruction set architecture (ISA) of a GPU. Accel-sim [24] is another example of an instruction-level GPU binary simulator. We note that such fine-grained, cycle-accurate or instruction-level simulators take a long time to simulate a single GPU, making them prohibitive for multi-GPU simulations. Also, ISA-level simulation is not relevant to distributed ML as the communication performance is dominated by network operations. Distributed ML Simulators. Multiple distributed ML simulators have been introduced to the research community. ATLAHS [42] models AI use cases, though it mainly looks at high-performance computing (HPC) or storage workloads. Calculon [21] provides detailed analytical modeling for LLM system co-design, and TrioSim [27] extrapolates single-GPU operator traces to large-scale multi-dimensional parallel execution. Maya [51] and Phanthora [37] reuse model code to easily capture and model workload information. SimAI [47] leverages parallelized discrete event simulation to improve scalability. Arcadia [11] scales further to support multi-job scheduling scenarios. Notably, these simulators don’t support arbitrarily customized collective algorithms and default to textbook algorithms. Also, none of the simulators capture CU-level, cache-line-sized GPU behaviors, crucial to modeling important control-path latencies.
7
Conclusion
We introduce ASTRA-sim 3.0, an update to the ASTRA-sim framework equipped with high-fidelity modeling capabilities. ASTRAsim 3.0 enables new design space exploration opportunities by capturing fine-grained control path and architectural details with
Table 1: ns-3 All-Reduce simulation result. Metric AllReduce Completion Time (𝜇s) Achieved Bus Bandwidth (Gbps) Min FCT (ns) Max FCT (ns) Avg FCT (ns) Standalone FCT (ns) Peak FCT Overhead (ns)
Clos Network 165.98 88.45 11,250 14,552 11,477 11,857 2,695
custom collective modeling. ASTRA-sim 3.0 also introduces InfraGraph, a common representation to capture the exact distributed ML infrastructure.
Acknowledgments We express our deepest gratitude to the original developers and authors of the ASTRA-sim infrastructure: Saeed Rashidi, Srinivas Sridharan, Sudarshan Srinivasan, and Taekyung Heo. AMD, the AMD Arrow logo, and combinations thereof are trademarks of Advanced Micro Devices, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies. AI Usage. Generative AI models (i.e., LLMs) are only used to check the writing quality and grammatical errors of the paper. All contents are solely created by the authors.
References AMD CDNA 4 Architecture. Accessed: 2026-04-30. [1] AMD. [n. d.]. https://www.amd.com/content/dam/amd/en/documents/instinct-techdocs/white-papers/amd-cdna-4-architecture-whitepaper.pdf [2] Zixian Cai, Zhengyang Liu, Saeed Maleki, Madan Musuvathi, Todd Mytkowicz, Jacob Nelson, and Olli Saarikivi. 2021. Synthesizing optimal collective algorithms. In Proc. undefined. 62–75. arXiv:2008.08708 [cs] doi:10.1145/3437801.3441620 [3] Francisco Caravaca, Ángel Cuevas, and Rubén Cuevas. 2025. From Prompts to Power: Measuring the Energy Footprint of LLM Inference. arXiv:2511.05597 [cs.AI] (2025). doi:10.48550/arXiv.2511.05597 [4] Xin Chen, Xiaoyang Wang, Ana Colacelli, Matt Lee, and Le Xie. 2025. Electricity Demand and Grid Impacts of AI Data Centers: Challenges and Prospects. (2025). doi:10.48550/arXiv.2509.07218 [5] Ziteng Chen, Xiaohe Hu, Menghao Zhang, Yanmin Jia, Yan Zhang, Mingjun Zhang, Da Liu, Fangzheng Jiao, Jun Chen, He Liu, Aohan Zeng, Shuaixing Duan, Ruya Gu, Yang Jing, Bowen Han, Jiahao Cao, Wei Chen, Wenqi Xie, Jinlong Hou, Yuan Cheng, Bohua Xu, Mingwei Xu, and Chunming Hu. 2025. An Efficient, Reliable and Observable Collective Communication Library in Large-scale GPU Training Clusters. arXiv:2510.00991 [cs] (2025). doi:10.48550/arXiv.2510.00991 [6] Jaehong Cho, Hyunmin Choi, Guseul Heo, and Jongse Park. 2026. LLMServingSim 2.0: A Unified Simulator for Heterogeneous and Disaggregated LLM Serving Infrastructure. arXiv:2602.23036 [cs] (2026). doi:10.48550/arXiv.2602.23036 [7] Jaehong Cho, Minsu Kim, Hyunmin Choi, Guseul Heo, and Jongse Park. 2024. LLMServingSim: A HW/SW Co-Simulation Infrastructure for LLM Inference Serving at Scale. In Proc. 2024 IEEE International Symposium on Workload Characterization (IISWC). 15–29. doi:10.1109/IISWC63097.2024.00012 [8] Sanghun Cho, Hyojun Son, and John Kim. 2023. Logical/Physical TopologyAware Collective Communication in Deep Learning Training. In Proc. 2023 IEEE International Symposium on High-Performance Computer Architecture (HPCA). 56–68. doi:10.1109/HPCA56546.2023.10071117 [9] Meghan Cowan, Saeed Maleki, Madanlal Musuvathi, Olli Saarikivi, and Yifan Xiong. 2023. MSCCLang: Microsoft collective communication language. In Proc. 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS). 502–514. doi:10.1145/3575693.3575724 [10] DeepSeek-AI. 2026. DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence. Accessed: 2026-05-01. https://huggingface.co/deepseekai/DeepSeek-V4-Pro/blob/main/DeepSeek_V4.pdf [11] Engineering at Meta. 2023. Arcadia: An end-to-end AI system performance simulator. Accessed: 2026-05-20. https://engineering.fb.com/2023/09/07/data-
ASTRA-sim 3.0: Next-Level Distributed Machine Learning Simulations via High-Fidelity GPU and Infrastructure Modeling
infrastructure/arcadia-end-to-end-ai-system-performance-simulator/ [12] Epoch AI. [n. d.]. Trends in Artificial Intelligence. Accessed: 2026-05-01. https: //epoch.ai/trends [13] William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961 [cs] (2022). doi:10.48550/arXiv.2101.03961 [14] gem5. [n. d.]. gem5: Interconnection network. Accessed: 2026-05-06. https://www. gem5.org/documentation/general_docs/ruby/interconnection-network/ [15] Prasun Gera, Hyojong Kim, Hyesoon Kim, Sunpyo Hong, Vinod George, and Chi-Keung Luk. 2018. Performance Characterisation and Simulation of Intel’s Integrated GPU Architecture. In Proc. 2018 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS). 139–148. doi:10.1109/ ISPASS.2018.00027 [16] Google Cloud. [n. d.]. What is a GPU & Its Importance for AI. Accessed: 2026-05-16. https://cloud.google.com/discover/gpu-for-ai [17] GPGPU-Sim. [n. d.]. GPGPU-Sim. Accessed: 2026-05-21. https://gpgpu-sim.org/ [18] Thomas R Henderson, Mathieu Lacage, and George F Riley. 2008. Network Simulations with the ns-3 Simulator. In Proc. Special Interest Group on Data Communication Conference (SIGCOMM). [19] Roger W. Hockney. 1994. The Communication Challenge for MPP: Intel Paragon and Meiko CS-2. Parallel Comput. 20 (1994), 389–398. https://api.semanticscholar. org/CorpusID:22986998 [20] Jiayi Huang, Pritam Majumder, Sungkeun Kim, Abdullah Muzahid, Ki Hwan Yum, and Eun Jung Kim. 2021. Communication algorithm-architecture co-design for distributed deep learning. In Proc. 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA). 181–194. doi:10.1109/ISCA52012. 2021.00023 [21] Mikhail Isaev, Nic McDonald, Larry Dennison, and Richard Vuduc. 2023. Calculon: A Methodology and Tool for High-Level Co-Design of Systems and Large Language Models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC ’23). Association for Computing Machinery, New York, NY, USA, Article 71, 14 pages. doi:10.1145/3581784.3607102 [22] Sylvain Jeaugey. 2019. Massively Scale Your Deep Learning Training with NCCL 2.4. Accessed: 2026-05-06. https://developer.nvidia.com/blog/massively-scaledeep-learning-training-nccl-2-4/ [23] Chelsea Maria John, Stepan Nassyr, Carolin Penke, and Andreas Herten. 2024. Performance and Power: Systematic Evaluation of AI Workloads on Accelerators with CARAML. In Proc. SC24-W: Workshops of the International Conference for High Performance Computing, Networking, Storage and Analysis. 1164–1176. doi:10.1109/SCW63240.2024.00158 [24] Mahmoud Khairy, Zhesheng Shen, Tor M. Aamodt, and Timothy G. Rogers. 2020. Accel-Sim: An Extensible Simulation Framework for Validated GPU Modeling. In Proc. 47th Annual International Symposium on Computer Architecture (ISCA). 473–486. doi:10.1109/ISCA45697.2020.00047 [25] Heehoon Kim, Junyeol Ryu, and Jaejin Lee. 2024. TCCL: Discovering better communication paths for pcie GPU clusters. In Proc. 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS). 999–1015. doi:10.1145/3620666.3651362 [26] Sabuj Laskar, Pranati Majhi, Sungkeun Kim, Farabi Mahmud, Abdullah Muzahid, and Eun Jung Kim. 2024. Enhancing collective communication in MCM accelerators for deep learning training. In Proc. 2024 IEEE International Symposium on High-Performance Computer Architecture (HPCA). doi:10.1109/HPCA57654.2024. 00069 [27] Ying Li, Yuhui Bao, Gongyu Wang, Xinxin Mei, Pranav Vaid, Anandaroop Ghosh, Adwait Jog, Darius Bunandar, Ajay Joshi, and Yifan Sun. 2025. TrioSim: A Lightweight Simulator for Large-Scale DNN Workloads on Multi-GPU Systems. In Proceedings of the 52nd Annual International Symposium on Computer Architecture (ISCA ’25). Association for Computing Machinery, New York, NY, USA, 1524–1538. doi:10.1145/3695053.3731082 [28] Xuting Liu, Behnaz Arzani, Siva Kesava Reddy Kakarla, Liangyu Zhao, Vincent Liu, Miguel Castro, Srikanth Kandula, and Luke Marshall. 2024. Rethinking Machine Learning Collective Communication as a Multi-Commodity Flow Problem. In Proc. Special Interest Group on Data Communication Conference (SIGCOMM). 16–37. doi:10.1145/3651890.3672249 [29] Jason Lowe-Power, Abdul Mutaal Ahmad, Ayaz Akram, Mohammad Alian, Rico Amslinger, Matteo Andreozzi, Adrià Armejach, Nils Asmussen, Brad Beckmann, Srikant Bharadwaj, Gabe Black, Gedare Bloom, Bobby R. Bruce, Daniel Rodrigues Carvalho, Jeronimo Castrillon, Lizhong Chen, Nicolas Derumigny, Stephan Diestelhorst, Wendy Elsasser, Carlos Escuin, Marjan Fariborz, Amin FarmahiniFarahani, Pouya Fotouhi, Ryan Gambord, Jayneel Gandhi, Dibakar Gope, Thomas Grass, Anthony Gutierrez, Bagus Hanindhito, Andreas Hansson, Swapnil Haria, Austin Harris, Timothy Hayes, Adrian Herrera, Matthew Horsnell, Syed Ali Raza Jafri, Radhika Jagtap, Hanhwi Jang, Reiley Jeyapaul, Timothy M. Jones, Matthias Jung, Subash Kannoth, Hamidreza Khaleghzadeh, Yuetsu Kodama, Tushar Krishna, Tommaso Marinelli, Christian Menard, Andrea Mondelli, Miquel Moreto, Tiago Mück, Omar Naji, Krishnendra Nathella, Hoa Nguyen, Nikos Nikoleris, Lena E. Olson, Marc Orr, Binh Pham, Pablo Prieto, Trivikram Reddy, Alec Roelke, Mahyar Samani, Andreas Sandberg, Javier Setoain, Boris Shingarov, Matthew D.
Sinclair, Tuan Ta, Rahul Thakur, Giacomo Travaglini, Michael Upton, Nilay Vaish, Ilias Vougioukas, William Wang, Zhengrong Wang, Norbert Wehn, Christian Weis, David A. Wood, Hongil Yoon, and Eder F. Zulian. 2020. The gem5 Simulator: Version 20.0+. arXiv:2007.03152 [cs] (2020). doi:10.48550/arXiv.2007.03152 [30] Junchao Ma, Dezun Dong, Cunlu Li, Ke Wu, and Liquan Xiao. 2021. PAARD: Proximity-Aware All-Reduce Communication for Dragonfly Networks. In Proc. 2021 IEEE Intl. Conf. on Parallel & Distributed Processing with Applications, Big Data & Cloud Computing, Sustainable Computing & Communications, Social Computing & Networking (ISPA/BDCloud/SocialCom/SustainCom). 255–262. doi:10.1109/ISPA-BDCloud-SocialCom-SustainCom52081.2021.00045 [31] Changhai Man, Joongun Park, Hanjiang Wu, Huan Xu, Srinivas Sridharan, and Tushar Krishna. 2025. STAGE: A Symbolic Tensor grAph GEnerator for distributed AI system co-design. arXiv:2511.10480 [cs] (2025). doi:10.48550/arXiv. 2511.10480 [32] Message Passing Interface Forum. 2025. MPI: A Message-Passing Interface Standard. Accessed: 2026-05-18. https://www.mpi-forum.org/docs/mpi-5.0/mpi50report.pdf [33] NVIDIA. [n. d.]. Pascal Tuning Guide - Pascal Tuning Guide 13.2 documentation. Accessed: 2026-05-20. https://docs.nvidia.com/cuda/pascal-tuning-guide/index. html [34] James O’Donnell and Casey Crownhart. [n. d.]. We did the math on AI’s energy footprint. Here’s the story you haven’t heard. MIT Technology Review. Accessed: 2026-05-14. https://www.technologyreview.com/2025/05/20/1116327/ai-energyusage-climate-footprint-big-tech/ [35] OpenAI. 2026. Introducing Stargate Norway. Accessed: 2026-05-14. https: //openai.com/index/introducing-stargate-norway/ [36] Ankur Patel. 2022. The Real Price of AI: Pre-Training Vs. Inference Costs. Accessed: 2026-05-13. https://www.ankursnewsletter.com/p/the-real-price-of-aipre-training [37] Jianxing Qin, Jingrong Chen, Xinhao Kong, Yongji Wu, Tianjun Yuan, Liang Luo, Zhaodong Wang, Ying Zhang, Tingjun Chen, Alvin R. Lebeck, and Danyang Zhuo. 2025. Phantora: Maximizing Code Reuse in Simulation-based Machine Learning System Performance Estimation. arXiv:2505.01616 [cs.DC] (2025). doi:10. 48550/arXiv.2505.01616 [38] Saeed Rashidi, Pallavi Shurpali, Srinivas Sridharan, Naader Hassani, Dheevatsa Mudigere, Krishnakumar Nair, Misha Smelyanski, and Tushar Krishna. 2020. Scalable Distributed Training of Recommendation Models: An ASTRA-SIM + NS3 case-study with TCP/IP transport. In Proc. 2020 IEEE Symposium on HighPerformance Interconnects (HOTI). 33–42. doi:10.1109/HOTI51249.2020.00020 [39] Saeed Rashidi, Srinivas Sridharan, Sudarshan Srinivasan, and Tushar Krishna. 2020. ASTRA-SIM: Enabling SW/HW Co-Design Exploration for Distributed DL Training Platforms. In Proc. 2020 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS). 81–92. doi:10.1109/ISPASS48437.2020. 00018 [40] Aashaka Shah, Vijay Chidambaram, Meghan Cowan, Saeed Maleki, Madan Musuvathi, Todd Mytkowicz, Jacob Nelson, Olli Saarikivi, and Rachee Singh. 2023. TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches. In Proc. USENIX Symposium on Networked Systems Design and Implementation (NSDI). [41] Aashaka Shah, Abhinav Jangda, Binyang Li, Caio Rocha, Changho Hwang, Jithin Jose, Madan Musuvathi, Olli Saarikivi, Peng Cheng, Qinghua Zhou, Roshan Dathathri, Saeed Maleki, and Ziyue Yang. 2025. MSCCL++: Rethinking GPU communication abstractions for cutting-edge AI applications. arXiv:2504.09014 [cs] (2025). doi:10.48550/arXiv.2504.09014 [42] Siyuan Shen, Tommaso Bonato, Zhiyi Hu, Pasquale Jordan, Tiancheng Chen, and Torsten Hoefler. 2025. ATLAHS: An Application-centric Network Simulator Toolchain for AI, HPC, and Distributed Storage. arXiv:2505.08936 [cs.DC] (2025). doi:10.48550/arXiv.2505.08936 [43] Srinivas Sridharan, Taekyung Heo, Louis Feng, Zhaodong Wang, Matt Bergeron, Wenyin Fu, Shengbao Zheng, Brian Coutinho, Saeed Rashidi, Changhai Man, and Tushar Krishna. 2023. Chakra: Advancing Performance Benchmarking and Co-design using Standardized Execution Traces. arXiv:2305.14516 [cs] (2023). [44] Rajeev Thakur, Rolf Rabenseifner, and William Gropp. 2005. Optimization of Collective Communication Operations in MPICH. The International Journal of High Performance Computing Applications 19, 1 (2005), 49–66. doi:10.1177/ 1094342005051521 [45] Guanhua Wang, Shivaram Venkataraman, Amar Phanishayee, Jorgen Thelin, Nikhil Devanur, and Ion Stoica. 2019. Blink: Fast and Generic Collectives for Distributed ML. In Proc. Conference on Systems and Machine Learning (SysML). [46] Haonan Wang, Xuxin Xiao, Mingyu Yan, Zhuoyuan Zhu, Dengke Han, Duo Wang, Wenming Li, Xiaochun Ye, Cunchen Hu, Hongyang Chen, and Guangyu Sun. 2025. A Systematic Characterization of LLM Inference on GPUs. arXiv:2512.01644 [cs] (2025). doi:10.48550/arXiv.2512.01644 [47] Xizheng Wang, Qingxu Li, Yichi Xu, Gang Lu, Dan Li, Li Chen, Heyang Zhou, Linkang Zheng, Sen Zhang, Yikai Zhu, Yang Liu, Pengcheng Zhang, Kun Qian, Kunling He, Jiaqi Gao, Ennan Zhai, Dennis Cai, and Binzhang Fu. 2025. SimAI:
Won et al.
unifying architecture design and performance tuning for large-scale large language model training with scalability and precision. In Proc. 22nd USENIX Symposium on Networked Systems Design and Implementation (NSDI). 541–558. [48] William Won, Midhilesh Elavazhagan, Sudarshan Srinivasan, Swati Gupta, and Tushar Krishna. 2024. TACOS: topology-aware collective algorithm synthesizer for distributed machine learning. In Proc. 57th IEEE/ACM International Symposium on Microarchitecture (MICRO). 856–870. doi:10.1109/MICRO61859.2024.00068 [49] William Won, Taekyung Heo, Saeed Rashidi, Srinivas Sridharan, Sudarshan Srinivasan, and Tushar Krishna. 2023. ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale. In Proc. 2023 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS). 283–294. doi:10.1109/ISPASS57527.2023.00035 [50] xAI. [n. d.]. Colossus: The World’s Largest AI Supercomputer. Accessed: 2026-05-06. https://x.ai/colossus [51] Srihas Yarlagadda, Amey Agrawal, Elton Pinto, Hakesh Darapaneni, Mitali Meratwal, Shivam Mittal, Pranavi Bajjuri, Srinivas Sridharan, and Alexey Tumanov. 2026. Maya: Optimizing Deep Learning Training Workloads using GPU Runtime Emulation. In Proc. 21st European Conference on Computer Systems (EuroSys). 1738–1758. doi:10.1145/3767295.3769366
[52] Yee Jiun Song and Kaushik Veeraraghavan. 2025. Meta’s Infrastructure Evolution and the Advent of AI. Accessed: 2026-05-06. https://engineering.fb.com/2025/09/ 29/data-infrastructure/metas-infrastructure-evolution-and-the-advent-of-ai/ [53] Jinsun Yoo, ChonLam Lao, Lianjie Cao, Bob Lantz, Minlan Yu, Tushar Krishna, and Puneet Sharma. 2025. Towards Easy and Realistic Network Infrastructure Testing for Large-scale Machine Learning. arXiv:2504.20854 [cs] (2025). doi:10. 48550/arXiv.2504.20854 [54] Jinsun Yoo, William Won, Meghan Cowan, Nan Jiang, Benjamin Klenk, Srinivas Sridharan, and Tushar Krishna. 2024. Towards a Standardized Representation for Deep Learning Collective Algorithms. In Proc. 2024 IEEE Symposium on HighPerformance Interconnects (HOTI). 33–36. doi:10.1109/HOTI63208.2024.00017 [55] Haidong Zhao and Nikolaos Georgantas. 2025. ML Inference Scheduling with Predictable Latency. In Proc. Proceedings of the Middleware for Autonomous AIoT Systems in the Computing Continuum (MAIoT). 25–30. doi:10.1145/3774901.3778066 [56] Liangyu Zhao, Saeed Maleki, Yuanhong Wang, Zezhou Wang, Ziyue Yang, Hossein Pourreza, and Arvind Krishnamurthy. 2026. ForestColl: ThroughputOptimal Collective Communications on Heterogeneous Network Fabrics. In Proc. USENIX Symposium on Networked Systems Design and Implementation (NSDI). arXiv:2402.06787