Starscream: Reimagining LLM Sharding on AMD Instinct GPUs with GPU Partitioning

Starscream: Reimagining LLM Sharding on AMD Instinct GPUs with GPU Partitioning Sumanth Gudaparthi† , Smeet Chheda† , and Ganesh Dasika

arXiv:2611.15951v1 [cs.DC] 11 Nov 2026

AMD Research and Advanced Development {sumanth.gudaparthi, smeet.chheda, ganesh.dasika}@amd.com

Abstract—As GPU accelerators grow larger to meet rising compute demands, they exhibit increasingly pronounced nonuniform memory access (NUMA) characteristics that challenge existing LLM frameworks. Chiplet-based architectures provide a practical path to larger GPUs while mitigating manufacturing yield constraints. Modern GPUs such as AMD Instinct MI300X employs such a multi-die design, and further offer GPU partitioning1 capabilities—partitioning a single physical GPU into multiple logical devices with NUMA-aware memory domains— promising faster kernel launches, and reduced memory access latencies with improved locality. However, current parallelization strategies, designed for homogeneous GPU topologies, when applied under partitioning, extend its unidimensional sharding causing significant performance degradation. Three bottlenecks emerge: (1) memory redundancy from replicated parameters or KV caches; (2) arithmetic intensity collapse as matrices become skinny and memory-bound; (3) communication overhead from flat collectives ignoring the 10× intra-GPU versus inter-GPU bandwidth difference. We present Starscream, a rethinking of LLM sharding for the chiplet era. Rather than stretching interGPU sharding thinner, Starscream adds intra-GPU sharding along an orthogonal dimension, grounded in three foundational principles: kernel shape preservation, memory redundancy elimination, and communication overhead minimization. These principles form the basis for extending any sharding strategy into the partitioning regime. We demonstrate their effectiveness using tensor parallelism as a case study, yielding sequence-length sharded attention eliminating KV replication, two-dimensional GEMM sharding maintaining arithmetic intensity, hierarchical collectives exploiting heterogeneous links, and computationcommunication fusion via symmetric memory. On AMD MI300X and MI355X, exploiting partitioning capabilities yields up to 4× speedup for attention and 1.43× for fused GEMMs compared to chiplet-oblivious baselines, demonstrating the effectiveness of chiplet-aware designs. Starscream’s principles generalize across sharding strategies and future architectures.

I. I NTRODUCTION The computational demands of large language models (LLMs) have driven an unprecedented expansion in GPU capabilities. Modern accelerators such as AMD’s Instinct MI300X offers a peak FP8 throughput of over 5 petaflops and 192 GB of HBM3 memory, while the following generation MI350 Series GPUs offers even greater performance through architectural refinements and superior memory bandwidth and capacity [7]. This trajectory shows no signs of slowing: as foundation models grow from hundreds of billions to trillions of parameters, hardware vendors continue to scale accelerator designs to meet demand [44]. † These authors contributed equally to this work. 1 not to be confused with sharding: LLM model split across multiple devices

(A)

(B)

Fig. 1. (A) AMD Instinct MI300X architecture. (B) Normalized selfattention time for Llama3-70B, Llama3-405B, and GPT-OSS120B on MI300X. Without partitioning: Purple line. With partitioning: red river is using naive sharding; blue river is using Starscream-enabled sharding. Bold colored lines show geomean speedup.

However, manufacturing economics impose fundamental constraints on monolithic die sizes [13], [14]. Larger dies suffer from exponentially decreasing yields as the probability of a defect appearing on any given die increases with area. Chipletbased architectures have emerged as the industry’s response to this challenge [37]. By decomposing a large monolithic design into multiple smaller dies interconnected through advanced packaging technologies, vendors can maintain high yields while still delivering aggregate compute density that would be hard with a single die. AMD’s MI300X GPU exemplifies this approach: the MI300X integrates eight Accelerated Compute Dies (XCDs) on top of 4 I/O Dies (IODs) each connected to a subset of HBMs, all connected through a sophisticated onpackage interconnect fabric. Future generations might expand this chiplet strategy further, potentially incorporating even more compute dies per package. Alongside compute scaling, memory bandwidth has evolved in lockstep. HBM has progressed with each generation, delivering substantial improvements in both capacity and bandwidth [8]. The MI300X, for instance, provides up to 5.3 TB/s of aggregate HBM bandwidth, a figure that will only increase with subsequent generations (e.g., up to 8 TB/s in MI350 Series GPUs). Yet these impressive bandwidth figures mask a critical architectural reality: this bandwidth is not uniformly accessible from all compute units. As dies grow larger or adopt

multi-chiplet designs, the physical distance between compute units and HBM controllers varies significantly—creating disparities between near and far memory access. This nonuniform memory access (NUMA) characteristic, long familiar in multi-socket CPU systems, is now becoming a first-class architectural concern in modern GPUs. AMD Instinct GPUs provide explicit mechanisms to exploit this inherent NUMA topology. Compute Partitioning modes logically divide the compute units (CUs) within a single physical GPU into independent virtual GPUs: DPX for 2-way, QPX for 4-way, and CPX for per-XCD granularity (8-way on MI300X). Each virtual GPU is capable of running separate workloads with dedicated compute resources. Complementing compute partitioning, the Memory Partitioning (NPS) mode subdivides the HBM address space into NUMA domains. In NPS4 mode on MI300X, each NUMA domain encompasses the memory attached to a single IOD, enabling workloads to achieve higher effective bandwidth by localizing memory accesses to nearby controllers [6]. We will refer to these capabilities collectively as partitioning. Figure 1A illustrates this chiplet-based architecture. At peak isolation, these partitioning modes unlock significant architectural benefits. Operating MI300X in CPX+NPS4 mode (8 virtual GPUs mapped across 4 NUMA domains) with no cross-die communication, provides substantial advantages over the default SPX+NPS1 configuration [10], [11]: • lower kernel launch latency: Avoiding L2 cache flushes between kernel runs, formerly required to preserve coherence across XCDs. • over 20% faster computations: HBM accesses limited to the local die reduce signal travel distance and eliminate contention on shared resources [10], [11]. • Reduced IOD power consumption: Lower cross-die traffic decreases power dissipation in the IODs. However, peak isolation is unrealistic for modern LLMs. Model parameters and KV caches exceed single-partition capacity, necessitating sharding across multiple devices, and how we shard determines what overheads get exposed. Current frameworks—both inference (vLLM [30], TensorRT-LLM [40]) and training (Pytorch [46], [56], Megatron-LM [51])—see virtual GPUs as independent devices, unaware of the partitioned topology. Total node resources remain unchanged, yet existing frameworks view virtual GPUs as additional devices. Naively extending the intra-node sharding strategies to the higher device count, though intuitive, exposes three fundamental bottlenecks: (1) reduced memory capacity from redundant memory replication, (2) degraded arithmetic intensity as matrices become increasingly narrow, and (3) communication overhead from flat collectives that ignore the bandwidth hierarchy between inter-die and inter-GPU links (inter-die links offer 10× inter-GPU bandwidth [6], [12]). The reason for this drop is, established strategies deliver performance by sharding each tensor along a single dimension to minimize communication—tensor parallelism (TP) over weight matrices [51], sequence parallelism (SP)

over inputs [1], [28], or FSDP for parameters and optimizer states [56]. Beyond a certain parallelism degree, however, this unidimensional sharding exposes memory, compute, and communication bottlenecks, degrading performance. These limitations manifest directly in partitioning configurations. The net effect: partitioning exposes a new level of hierarchy that unlocks significant performance benefits. However, without such awareness, these advantages remain latent. Figure 1B (red-river) shows an extreme example: 8 KV-head attention on 8 GPUs (64 in CPX), where the penalty is severe. We present Starscream, a rethinking of LLM execution for partitioned GPUs. The core insight: leverage the coarsegrained control over data placement and compute mapping that partitioning modes provide to maximize data locality and minimize cross-die accesses. Rather than treating partitioned GPUs as simply “more GPUs,” Starscream recognizes that partitioning exposes a new level of hierarchy, and composes existing parallelism strategies hierarchically—applying each at the level of the topology where it is effective to extract performance from the same compute resources. This principle yields a systematic set of techniques: • Eliminating memory pressure: Hierarchical sharding avoids replicating data structures that naive sharding duplicates—KV caches in TP inference, parameters in FSDP training—preserving memory capacity and bandwidth. • High Arithmetic Intensity GEMMs: 2D sharding maintains arithmetic intensity even as parallelism increases, keeping operations compute-bound rather than memory-bound. • Hierarchical collectives: Collectives structured as twolevel operations: fast intra-GPU operations over USR links, followed by inter-GPU synchronization on XGMI. • Computation-communication fusion: Overlaps kernel execution (attention, GEMMs, RMSNorm) with inter-XCD data movement, hiding collective latency behind computation. We implement and evaluate Starscream on AMD Instinct MI300X GPUs, demonstrating that chiplet-aware design transforms GPU partitioning into a significant advantage. Our NUMA-aware paged-attention kernels achieve up to 4× speedup in CPX/NPS4 mode over SPX/NPS1. Our fused GEMM kernels with inter-XCD peer-mapping deliver 1.43× faster execution. These per-kernel improvements compound to yield substantial end-to-end throughput gains across diverse model architectures and batch sizes. Beyond raw performance, Starscream’s hierarchical sharding principles generalize across sharding strategies, offering a blueprint for exploiting NUMA topologies in future chiplet-based accelerators. The remainder of this paper is organized as follows. Section II explains how partitioning yields performance benefits. Section III presents the challenges that current LLM frameworks face under partitioning, identifying the root causes for performance degradation. Section IV describes the proposed Starscream hierarchical sharding strategies. Section V details our experimental methodology and evaluation results demonstrating the effectiveness of our approach. Section VI shares the related works and we conclude in Section VII.

II. PARTITIONING ON AMD I NSTINCT MI300X GPU S This section describes the AMD MI300X architecture and its compute/memory partitioning modes. While we focus on MI300X, the principles apply to most Instinct GPUs broadly. A. MI300X Chiplet Architecture Figure 1A illustrates the MI300X architecture. At its core are four I/O Dies (IODs) connected via a ring-style Ultra Short Reach (USR) interconnect. The USR interconnect bandwidth is significantly higher than the XGMI bandwidth [9], [12]. Each IOD hosts two Accelerator Compute Dies (XCDs) and connects to HBM memory slices. Within each XCD, 38 Compute Units (CUs) execute SIMD operations independently. Each XCD has a dedicated L2 cache slice, and a shared Infinity Cache spans each IOD, acting as the last-level cache (LLC). B. Baseline Operation: SPX NPS1

remote XCDs are unaware of. Maintaining coherence requires flushing the L2 cache before each kernel launch, increasing launch latency by a non-trivial factor on MI300X. IOD power consumption. Every remote HBM access requires the data to first travel from HBM through the IOD before reaching the USR interconnect. The IODs are thus actively involved in all cross-die data movement. Higher remote traffic translates to higher IOD power consumption, leaving less power budget for XCD compute operations. The GPU operates within a fixed power envelope, and power consumed by the IODs is power unavailable to the XCDs. With a reduced power budget, XCDs can only clock so high before hitting their power limit—resulting in lower operating frequencies and corresponding performance degradation. D. Compute and Memory Partitioning

In the default Single Partition Accelerator (SPX) mode with unified memory (NPS1), the GPU presents the illusion of a single monolithic device. The L2 cache appears shared across all XCDs, and HBM is accessible from any compute unit. Memory allocation. Assuming a uniform memory distribution, pages would be allocated to HBM slices in round-robin fashion across all dies. The first page goes to HBM slice 0 on IOD 0, the second to HBM slice 1 on IOD 0, the third to HBM slice 0 on IOD 1, and so forth. Data access paths. When a CU requires data, the access latency depends on where the data resides. If the data is in a local HBM slice—one connected to the same IOD as the requesting XCD—it is a direct read. However, if the data resides in a remote HBM slice on a different IOD, the data must traverse the USR ring: it is first read from the remote HBM to its local IOD, then transferred across one or more USR hops to the requesting IOD, and finally delivered to the XCD. For the farthest memory (two hops away), this path is significantly longer than local access.

AMD Instinct GPUs can partition a single physical GPU into multiple compute and memory domains. We focus on the most aggressive configuration: CPX NPS4. Compute Partitioning (CPX) mode partitions compute resources at XCD granularity. Each XCD becomes an independent virtual GPU, receiving 1/8th of total HBM capacity. The XCD’s local L2 slice now becomes its private L2 cache. Memory Partitioning (NPS4). NPS4 mode partitions memory into four NUMA domains, one per IOD. Each IOD’s HBM slices constitute a separate NUMA domain. With CPX NPS4, each XCD’s 1/8th memory allocation is exclusive to its local NUMA domain—the HBM slices connected to its IOD. Since each IOD hosts two XCDs, half of the IOD’s HBM capacity is allocated to each XCD. Bandwidth implications. In NPS1, each XCD can theoretically tap the aggregate bandwidth of all HBM slices across all IODs. In NPS4 on MI300X, each XCD can only access bandwidth from the HBM slices on the same IOD. This represents a trade-off: reduced aggregate bandwidth accessibility in exchange for guaranteed local bandwidth without contention.

C. Challenges in SPX NPS1 Mode

E. Benefits of CPX NPS4 Mode

Three challenges emerge from this unified memory model: Remote memory access overhead. Since data is uniformly distributed across all HBM slices, a significant fraction of memory accesses target remote IODs, incurring the multihop latency described above. The problem compounds under concurrent access. When multiple XCDs simultaneously request data from different remote IODs, their traffic converges on the shared USR fabric, creating contention. The ring topology means that traffic between distant IODs must traverse intermediate links, competing with other flows. Even though HBM offers substantial aggregate bandwidth, the USR network becomes the bottleneck—effective bandwidth is limited by the slowest, most congested link. Combined, these two factors—longer data paths and USR contention—can account for over 20% of total compute time. L2 cache coherence overhead. Because L2 is sliced across XCDs, coherence must be maintained across dies. After a kernel completes, one XCD may hold dirty cache lines that

Operating in CPX NPS4 mode in perfect peak isolation eliminates the three challenges identified above: No L2 coherence overhead. Each XCD has its own private L2 cache and operates as an independent GPU. There is no cross-die L2 sharing, no coherence concerns, and consequently no need for L2 flushes between kernel launches. No remote memory traffic. In NPS4 mode, all memory requests from an XCD are routed exclusively to HBM slices within the same IOD. There is no cross-die memory traffic, eliminating USR contention and remote access latency. Reduced IOD power consumption. Without cross-die data movement, the IODs consume less power, preserving power budget for XCD compute operations and enabling higher sustained frequencies. These benefits motivate our focus on CPX NPS4 as the operating mode for GPU partitioning. However, as discussed in Section I, realizing these benefits requires rethinking parallelization strategies—a challenge we address with Starscream.

III. A NATOMY OF NA ÏVE S HARDING WITH PARTITIONING We motivate the challenges of partitioning through a concrete example: deploying a Grouped Query Attention (GQA) [4] model using tensor parallelism [51] on an illustrative node with GPUs supporting compute and memory partitioning. This example exposes the three fundamental bottlenecks discussed earlier: data replication, arithmetic intensity degradation, and collective communication overhead; that cause naive sharding to degrade performance despite its architectural advantages. A. Illustrative System Setup and Model Characteristics Hardware Configuration. We consider a node equipped with 4 physical GPUs interconnected via AMD’s Infinity Fabric (XGMI) [5], shown in Figure 2A. Each GPU contains 2 IODs and 2 XCDs (one per IOD), linked by high-bandwidth USR connections. In the baseline SPX/NPS1 mode, software frameworks perceive 4 monolithic GPUs. Enabling DPX/NPS2 mode partitions each physical GPU into 2 virtual GPUs, exposing 8 logical devices to the framework, each with half the memory capacity and bandwidth. Model Configuration. Our illustrative LLM layer uses Grouped-Query Attention (GQA) with 8 attention heads and 4 KV heads, where each KV head is shared by 2 query heads. The layer’s computation comprises three kernel categories: (1) Attention kernels perform self-attention by querying the KV cache; (2) GEMM kernels handle input projection (QKV generation), output projection, and two MLP layers (FC1, FC2); (3) Collective kernels synchronize partial results across GPUs via all-reduce operations. While modern LLMs include additional kernel types, we focus on those that dominate execution time in this illustration. Figure 2B illustrates these components and their tensor-parallel sharding under the baseline 4-GPU configuration.

(A) Illustrative node

(B) Illustrative LLM model Fig. 2. Illustrative (a) partitioned node and (b) LLM model.

B. Attention Kernels: KV Cache Replication Tensor parallelism distributes attention computation via head-parallel sharding: each GPU receives a subset of query heads along with their corresponding KV heads. In the baseline 4-GPU configuration (Figure 2A), each GPU handles 2 query heads and 1 KV head (top half of Figure 3A). Since both

query heads share the same KV head for their self-attention computation, and all three reside on the same GPU, this sharding requires no cross-GPU communication and incurs no redundant memory overhead. Partitioning disrupts this clean mapping. With 8 virtual GPUs, head-parallel sharding assigns only 1 query head per device, but only 4 KV heads exist to serve 8 devices. If we naively increase tensor parallelism from 4 to 8 to match the new device count, current frameworks (vLLM [30], MegatronLM [51]) resolve this mismatch by replicating KV heads: each virtual GPU receives a full copy of its required KV head (bottom half of Figure 3A). Consequently, each physical GPU now stores the same query heads as before (1 per XCD), but maintains 2 redundant copies of the same KV head across its two XCDs. This replication carries severe consequences. The KV cache already consumes 3× more memory than model parameters in long-context scenarios [45]. Duplicating it across virtual GPUs reduces effective memory capacity per physical GPU, directly limiting maximum batch size and supported context length. Worse, both XCDs now read from separate local copies rather than sharing cached data, halving effective memory bandwidth. On a real MI300X with 8-way partitioning, this replication amplifies to 8×, significantly reducing both available memory and bandwidth. Since attention is memorybandwidth bound during autoregressive decode, execution time increases proportionally—and the system hits out-of-memory far sooner than baseline, as shown in Figure 1B. C. GEMM Kernels: Arithmetic Intensity Collapse Tensor parallelism [51] shards GEMM operations using two complementary strategies.2 Consider a GEMM computing CM ×N = AM ×K · WK×N distributed across P GPUs: • ColumnParallelLinear shards the weight matrix along the column dimension (i.e., N ). Each GPU stores WK×N/P and receives a full copy of the input AM ×K , then independently computes a non-overlapping output slice CM ×N/P . FC1, the first MLP layer, uses this strategy (Figure 2B). • RowParallelLinear shards along the row dimension (i.e., K). Each GPU receives input AM ×K/P —matching the sharded output from FC1—and stores a correspondingly sharded weight WK/P ×N . Each GPU computes a full CM ×N result, but because K is partial, these must be aggregated via all-reduce. This strategy is typically used in the second MLP layer (FC2) in Figure 2B. We focus on RowParallelLinear, illustrated in the last layer of Figure 2B. In the baseline 4-GPU configuration, each GPU receives a K/4 × N weight shard. Under partitioning with 8 virtual GPUs (Figure 3B), each receives K/8×N . On a typical MI300X node with 8 physical GPUs operating in CPX mode, this degrades further to K/64 × N per virtual GPU. The problem is that these sharding strategies employ unidimensional sharding: sharding typically occurs along the same dimension (K for RowParallelLinear, N for ColumnParallelLinear) regardless of parallelism degree. As P grows, this 2 Terms ColumnParallelLinear, RowParallelLinear used from vLLM [30].

(A)

(B)

(C)

Fig. 3. Illustrative GPU sharding for (A) GQA and (B) RowParallelLinear. (C) Arithmetic intensity: 8 GPUs operating in SPX versus CPX.

produces increasingly skinny, elongated matrices that are inefficient to compute. A K/64 × N shard has dramatically lower arithmetic intensity than the original K × N matrix, even though total computation and hardware resources remain constant. Operations compute-bound at P = 1 might become memory-bound at P = 64, underutilizing available FLOPS. Figure 3C quantifies this drop for 8-way versus 64-way sharding across various models. D. Collective Kernels: Quadratic Communication Growth RowParallelLinear requires all-reduce to aggregate partial outputs across GPUs. In the baseline 4-GPU configuration (Figure 2B), each GPU holds an M × N partial result and exchanges data with 3 peers. Assuming ring all-reduce, each GPU sends and receives (P − 1) · M · N elements. Partitioning dramatically increases the peer count. With 8 virtual GPUs, each must now communicate with 7 peers, sending (8−1)·M ·N elements. Although total computational work remains constant, communication volume grows with parallelism degree: doubling P from 4 to 8 increases perGPU data movement from 3 · M · N to 7 · M · N —a 2.3× increase for a 2× increase in parallelism. For 64-way parallelism on an 8-GPU MI300X node, each virtual GPU exchanges data with 63 peers, resulting in 21× more communication than the baseline. While more advanced all-reduce algorithms can mitigate this overhead, the fundamental scaling problem persists. Moreover, naive sharding treats all virtual GPUs uniformly, ignoring the heterogeneous link topology. Communication between virtual GPUs on the same physical die (via USR links) enjoys 10× higher bandwidth than cross-die communication (via XGMI) [6]. Current frameworks perform flat all-reduces that fail to exploit this hierarchy: all transfers proceed at the rate of the slowest link (XGMI), leaving the faster intra-die bandwidth (USR) underutilized. E. Training: FSDP2 Exhibits Identical Pathologies While the preceding analysis focused on tensor-parallel inference, training workloads using Fully Sharded Data Parallel (FSDP) [46], [56] exhibit identical pathologies. FSDP shards model parameters, gradients, and optimizer states across GPUs while operating on a different batch. During forward and backward passes, each GPU all-gathers layer parameters before computation and reduce-scatters gradients after, respectively.

In the baseline 4-GPU configuration, frameworks overlap communication with computation: while GPU i computes layer ℓ, it concurrently all-gathers parameters for layer ℓ + 1. Similarly, during backward passes, reduce-scatter for layer ℓ overlaps with gradient computation for layer ℓ − 1. This pipelining masks communication latency, though some exposure remains for large models or small micro-batches. Partitioning breaks this balance. Each virtual GPU now allgathers the entire layer’s parameters to its local memory— which has half the capacity and bandwidth of the original physical GPU. For an 8-way partitioned MI300X, this means each XCD must reconstruct full layer parameters into 1/8th the memory with 1/8th the bandwidth. The collective scope also increases from 3 peers (4-GPU baseline) to 7 peers (8-way sharding), amplifying communication overhead. The arithmetic intensity problem also persists. As P grows, the data batch per P become increasingly skinny, reducing perGPU arithmetic intensity and underutilizing compute resources during local forward/backward computations. Summary: Universal Bottlenecks The bottlenecks identified—memory capacity reduction, collective communication overhead, and arithmetic intensity degradation—are universal across sharding strategies and workload types. Whether applying tensor parallelism or FSDP2, whether performing inference (prefill or decode) or training (forward or backward), naive sharding on partitioned GPUs amplifies the same fundamental issues: • Memory pressure: Fractional capacity/bandwidth limits batch size, context length, and model scale. • Communication inflation: Data movement scales with partitioning; flat collectives skip link hierarchy. • Compute underutilization: Unidimensional sharding yeilds skinny, low arithmetic intensity matrices. These pathologies compound to overwhelm partitioning’s architectural benefits (significantly faster kernel launch, low latency accesses, and higher power window), causing significant end-to-end performance degradation compared to the SPX/NPS1 baseline.

IV. S TARSCREAM Partitioning fundamentally transforms the hardware substrate. A single physical GPU becomes several smaller virtual GPUs—each with reduced memory capacity, diminished bandwidth, a fraction of the L2 cache, and fewer compute units. These are not minor perturbations; they reshape the performance characteristics that every sharding strategy was tuned to exploit. Every parallelism strategy carries inherent overheads that amplify or remain hidden depending on hardware characteristics. When partitioning changes these characteristics, bottlenecks emerge—affecting enough kernels to degrade end-to-end execution. Section III quantified these pathologies. Starscream addresses this gap. By systematically analyzing the architectural changes that partitioning introduces, Starscream defines fundamental principles that capture the benefits of partitioning while paying minimal cost. The key insight is deceptively simple: do not extend existing sharding deeper—shard orthogonally instead. We present Starscream in two parts. First, we introduce the orthogonal sharding principle—the theoretical foundation that explains why naive sharding fails and what alternative succeeds. Second, we apply this principle to tensor-parallel inference, detailing how orthogonal sharding transforms selfattention, GEMMs, and collectives on AMD hardware. A. The Orthogonal Sharding Principle The key insight behind Starscream is multi-dimensional sharding. When workloads are sharded across physical GPUs in a node along a given dimension, switching to partitioning should not always extend that same sharding deeper— this is precisely what causes naive sharding to fail. Instead, partitioning should retain the work per physical GPU exactly as the ideal sharding for that GPU topology suggests, while introducing sharding along an orthogonal dimension across the virtual GPUs within each physical GPU. This orthogonal dimension is chosen to simultaneously optimize three critical tradeoff points, as illustrated in Figure 4: 1) Kernel shapes: The orthogonal shard must produce matrix tiles that maintain high arithmetic intensity, avoiding the skinny elongated shapes that shift operations from compute-bound to memory-bound. 2) Memory redundancy: The orthogonal shard must avoid replicating tensors (such as KV caches) that would consume precious memory capacity and bandwidth. 3) Communication overhead: The orthogonal shard must minimize cross-die data movement time by exploiting the hierarchical bandwidth structure—fast intra-GPU links versus slower inter-GPU links. The choice of which orthogonal dimension to shard along under partitioning is determined by analyzing these three tradeoff points for the specific inter-GPU sharding strategy in use. Different parallelism strategies—tensor parallelism (TP), fully sharded data parallelism (FSDP), sequence parallelism (SP)—each shard along different dimensions across GPUs, and

consequently each requires a different orthogonal dimension for effective performance. Crucially, orthogonal sharding is compositional: Starscream does not replace existing parallelism strategies—it augments them. A framework can continue using tensor parallelism, FSDP, or any combination across physical GPUs exactly as before; Starscream adds NUMA-aware partitioning at the intra-GPU level. This design philosophy transforms GPU partitioning into an opportunity to extract additional performance from the same compute resources.

Fig. 4. The three tradeoffs guiding orthogonal sharding.

B. Scope: Tensor Parallelism as a Case Study While Starscream provides a comprehensive framework with solutions for each sharding strategy (TP, FSDP, SP, PP, and their combinations), the underlying principles are universal. Given the limited space in this paper, we focus on tensor parallelism for LLM inference as a detailed case study. This choice is deliberate: • Tensor parallelism is the dominant strategy for largemodel inference, where latency-sensitive decode phases make partitioning’s locality benefits most attractive. • The three bottlenecks—memory replication, GEMM intensity collapse, and collective overhead—manifest most acutely under tensor parallelism with Grouped Query Attention models [4]. • The orthogonal sharding solutions for tensor parallelism directly demonstrate how analyzing the three tradeoff points guides dimension selection. Since the principles are fundamental, the insights from this case study transfer directly to other sharding strategies and attention variants ( MQA [50], MLA [18], CHAI [3]). For example, extending FSDP to partitioning scope would require all-gathering parameters into smaller memory while increasing microbatch counts, both compound to burden the memory capacity limiting the max model size supported. The orthogonal approach instead shards along dimensions (hidden, sequence) that avoid expanding all-gather scope. The analysis framework remains identical; only the specific dimension choices differ. C. Self-Attention with Sequence-Length Sharding As established in Section III, tensor-parallel self-attention results in KV cache replication when the number of KV heads is fewer than the number of virtual GPUs—a common scenario with 64 virtual GPUs per node in CPX/NPS4 mode and

Fig. 5. Starscream sequence-length sharding: two self-attention variants.

modern architectures using GQA or MLA. Since memory capacity does not scale with partitioning, naive sharding creates multiple redundant copies of the same KV cache, degrading both capacity and bandwidth. a) Orthogonal Sharding Strategy: Instead of replicating the KV cache, Starscream shards it across compute partitions along the sequence length dimension. Each virtual GPU holds a subset of tokens, allocated in round-robin fashion: the first token’s KV entry resides on XCD0, the second on XCD1, and so on. Figure 5 illustrates this sharding. Each XCD retains its assigned query head from inter-GPU tensor parallelism, but rather than duplicating the KV cache, XCDs share a single distributed copy. This eliminates memory redundancy entirely. b) Computation with Distributed KV Cache: Since selfattention requires each query to attend over all KV tokens, the distributed cache necessitates coordination. Two approaches are possible: 1) KV all-gather: Before computing each layer, all-gather the KV cache entries across XCDs. Each XCD then holds a complete copy, computes full self-attention locally, and evicts the gathered entries afterward—analogous to how FSDP all-gathers parameters per layer. 2) Query all-gather with partial attention: Keep KV entries in place and all-gather the query heads instead. Each XCD then holds all query heads but only its local KV subset, computing partial self-attention. These partials are then exchanged via a collective call, followed by a lightweight aggregation kernel. c) Decode versus Prefill: The choice between these approaches depends on the inference phase. During decode, the autoregressive nature produces exactly one query token per iteration, while the KV cache holds all tokens observed thus far. Moving queries is far cheaper than moving the entire KV cache—we select query all-gather. During prefill, all input tokens enter simultaneously, so query and KV token counts are comparable. However, GQA models have far more query heads than KV heads (e.g., Llama-3-405B: 128 query heads, 8 KV heads). Since KV head volume is 16× smaller, we select KV all-gather for prefill. Notably, KV all-gather enables full local attention, eliminating the aggregation kernel. d) Effectiveness: This approach dramatically reduces inter-die data movement. Consider Llama-3-405B on MI300X in CPX/NPS4 (64 virtual GPUs) with 128K context. During decode, transferring only query heads and aggregation meta-

data accounts for less than 0.1% of the data consumed by the attention computation. During prefill, KV all-gather accounts for only 4.4% of total data—a 16× reduction from query allgather due to the GQA head ratio. Compared to the baseline where 75% of memory accesses were remote (Section II), Starscream reduces cross-die traffic to under 5%. Beyond reduced latency, sequence-length sharding delivers two additional benefits. First, since each KV entry is read exactly once per physical GPU (on whichever XCD owns it), memory bandwidth is no longer wasted reading duplicate copies. Second, inter-XCD traffic occurs via controlled collective calls rather than scattered remote accesses, avoiding contention and maximizing HBM bandwidth utilization. D. Hierarchical Collective Primitives As discussed in Sections II and III, conventional collective calls treat all links uniformly, ignoring the bandwidth difference between much faster USR links and XGMI links. When naive sharding splits workloads across all XCDs in a node, partials must be aggregated across the entire topology. A flat all-reduce spanning all 64 virtual GPUs in MI300X bottlenecks on the slowest link—XGMI—leaving the faster USR fabric underutilized. a) Three-Phase Hierarchical All-Reduce: Starscream restructures all-reduce as a three-phase operation that decouples intra-GPU and inter-GPU communication. We illustrate using the 4-GPU, 2-XCD-per-GPU system from Figure 3, where each colored tile represents an M × N partial result. 1) Intra-GPU reduce-scatter (Phase 1): Each physical GPU performs a local reduce-scatter across its XCDs over USR links. As shown in Figure 6A, after this phase each XCD holds half the tile (shape M × N/2), containing the reduced result from both local XCDs. All physical GPUs execute this phase simultaneously with no crossGPU dependency. 2) Inter-GPU peer-XCD all-reduce (Phase 2): XCDs at the same position across physical GPUs (all XCD0s, all XCD1s) perform an all-reduce over XGMI links. As shown in Figure 6B, after this phase each XCD holds the fully reduced result for its half of the tile—aggregated across all physical GPUs. 3) Intra-GPU all-gather (Phase 3): Each physical GPU performs a local all-gather across its XCDs over USR links. As shown in Figure 6C, this distributes the complete M ×N result to all XCDs. Like Phase 1, all physical GPUs execute in parallel. b) Benefits: This decomposition provides two key advantages. First, intra-GPU and inter-GPU communication are fully decoupled: the fast USR links handle Phases 1 and 3 without being throttled by slower XGMI traffic. Second, the inter-GPU all-reduce in Phase 2 operates over physical GPUs only (e.g., 8 participants), not virtual GPUs (64 participants), reducing XGMI traffic proportionally. Together, these ensure both network tiers operate at their respective peak bandwidths. Furthermore, because intra-GPU communication is isolated, we can tailor Phases 1 and 3 to exploit MI300X’s ring topol-

⇒

(A)

⇒

(B)

(C)

Fig. 6. The three-stage hierarchical all-reduce: (A) intra-GPU reduce-scatter, (B) inter-GPU peer-XCD all-reduce, (C) intra-GPU all-gather.

ogy without affecting the inter-GPU collective. This enables topology-aware optimizations—such as direction-optimized ring reduce-scatter—that would be impossible in a flat collective spanning heterogeneous links. E. Two-Dimensional GEMM Sharding As shown in Section III, extending unidimensional sharding into partitioning creates increasingly skinny matrices with degraded arithmetic intensity. If RowParallelLinear shards along K across physical GPUs, naive sharding splits further along K across XCDs, producing K/64 × N shards that are severely memory-bound. a) Orthogonal Sharding Strategy: Starscream retains the original inter-GPU sharding and selects the orthogonal dimension for intra-GPU sharding. If physical GPUs use RowParallelLinear (sharding along K), XCDs use ColumnParallelLinear (sharding along N ), and vice versa. This keeps tiles approximately square, preserving arithmetic intensity. b) Collective Interchange: Different sharding strategies require different collectives: RowParallelLinear produces partials requiring reduce-scatter at the output, while ColumnParallelLinear requires all-gather at the input to reconstruct the full activation. Crucially, these are interchangeable: for the same per-GPU GEMM work, both approaches produce identical output shapes—only the collective placement differs. Consider FC2 with RowParallelLinear across physical GPUs. Each physical GPU holds a K/4×N weight shard, and each XCD receives input of shape M ×K/8 from the previous layer. Two options exist for intra-GPU sharding (Figure 7): 1) Extend RowParallelLinear: Shard weights further along K, giving each XCD a K/8 × N shard. Each XCD computes M ×N partials, followed by inter-XCD reducescatter, yielding M × N/2 output per XCD. 2) Apply ColumnParallelLinear: Shard weights along N , giving each XCD a K/4 × N/2 shard. First all-gather the input to M × K/4, then compute the GEMM, yielding M × N/2 output per XCD. Both produce identical M × N/2 outputs per XCD. Option 1 creates skinny K/8 × N tiles; Option 2 creates balanced K/4 × N/2 tiles with 2× better arithmetic intensity. c) Applying to the Full LLM: In standard tensor parallelism, all-reduce boundaries occur between output projection and FC1, and between FC2 and the next layer’s input projection. Output projection and FC2 use RowParallelLinear; FC1 and input projection use ColumnParallelLinear. Replacing all-reduce with hierarchical collectives (Section IV-D), we

observe: RowParallelLinear GEMMs are followed by interXCD reduce-scatter, and ColumnParallelLinear GEMMs are preceded by inter-XCD all-gather. Using the collective interchange property, we can substitute any “GEMM → reduce-scatter” with “all-gather → GEMM”—converting RowParallelLinear to ColumnParallelLinear at the XCD level, or vice versa. This flexibility allows all GEMMs in an LLM layer to be 2D-sharded: orthogonal dimensions across physical GPUs and XCDs, maintaining balanced tiles and high arithmetic intensity throughout. F. Computation-Communication Overlap Even with hierarchical collectives and 2D sharding, discrete kernel boundaries impose three performance taxes [53]: (1) HBM write at compute kernel completion and read at collective kernel start, (2) bulk synchronization across all workgroups within each kernel, and (3) kernel launch overhead for the collective. These taxes compound across the many compute-collective pairs in an LLM layer. Starscream eliminates these taxes by fusing compute and collective operations. The key insight: if a kernel’s output must reside in a specific NUMA region for the subsequent kernel, we customize the compute kernel to write directly to that region as outputs are generated. Each thread block writes to the appropriate memory partition using symmetric memory [2], avoiding a separate collective kernel entirely. We apply this fusion throughout the LLM pipeline— GEMMs, attention, RMSNorm, and any operation requiring cross-XCD communication under our sharding strategy. Whether the collective precedes (all-gather) or follows (reduce-scatter) the computation, we use software pipelining,

(A)

(B)

Fig. 7. Orthogonal GEMM sharding: (A) Work allocated per GPU, (B) inter-XCD sharding: RowParallelLinear and ColumnParallelLinear.

Fig. 8. (A) Normalized Speedup, (B) memory bandwidth utilization, (C) Additional Speedup from compute-communication overlap for GPT-OSS120B, Llama3-70B and Llama3-405B models with varying batch size and context lengths.

prefetching, and careful data orchestration to overlap communication with compute, hiding whatever latency remains after hierarchical mapping. V. E VALUATION A. Methodology a) Hardware: We evaluate Starscream on AMD Instinct MI300X and MI355X GPUs. For MI300X, we compare CPX/NPS4 mode against baseline SPX/NPS1 mode. For MI355X, we compare DPX/NPS2 mode against SPX/NPS1. b) Implementation: Self-attention kernels build on vLLM’s [30] Triton-based unified-attention implementation (evaluated on vLLM-v1) and the HIP-based [31] pagedattention implementation (evaluated on vLLM-v0). GEMM kernels are implemented using the Stream-K [42] algorithm in Triton. Computation-communication fusion leverages PyTorch’s symmetric memory [2]. We plan to open-source all implementations upon acceptance. c) Workloads: For self-attention, we evaluate the decode phase for three representative model shapes—GPT-OSS-120B, Llama3-70B, and Llama3-405B—across context lengths from 16K to 128K tokens and batch sizes from 1 to 1024. For GEMM and collective kernels, we use over 200 representative shapes drawn from state-of-the-art LLMs including GPT-OSS120B, Llama3-70B, Llama3-405B, and others. We deliberately selected shapes representing both state-of-the-art GQA-based and MoE-based architectures to ensure broad coverage. B. Self-Attention Performance Figure 8A shows the speedup of self-attention kernels during the decode phase using Starscream’s sequence-length sharding (CPX Starscream) and naive sharding (CPX baseline)

in CPX/NPS4 mode relative to the SPX/NPS1 baseline. Figure 8B shows the corresponding HBM bandwidth utilization for the same configurations. a) Low-Latency Scenarios: At small batch sizes, Starscream achieves its highest speedups. In SPX/NPS1 mode, small batches create memory access hotspotting—requests concentrate on a subset of HBM slices, leaving aggregate bandwidth underutilized. Starscream’s sequence-length sharding distributes the KV cache evenly across NUMA domains, eliminating hotspots and saturating available bandwidth. This translates directly to improved throughput. b) High-Throughput Scenarios: As batch sizes increase, the performance gap between Starscream and SPX/NPS1 narrows. Larger batches naturally distribute memory accesses more uniformly, reducing hotspotting even in SPX mode. Figure 8B confirms this: SPX bandwidth utilization improves with batch size, approaching Starscream’s utilization at the largest batches (reflected in the smaller SPX-CPX gap). However, SPX/NPS1 at high batch sizes introduces a different bottleneck: USR link contention. With more data flowing across IODs, XCDs compete for shared interconnect bandwidth, reducing net effective bandwidth. Starscream sidesteps this by minimizing USR traffic—all-gathering the smaller tensor (queries or KV heads, depending on phase) keeps crossdie communication lightweight, sustaining performance gains even when hotspotting is no longer the dominant factor. c) Naive Sharding Fails: The CPX/NPS4 baseline without Starscream (naive sharding) performs consistently worse than SPX/NPS1, and degrades further as batch size increases. KV cache replication is the culprit: larger batches require more KV entries, and replicating them across XCDs consumes memory bandwidth bringing redundant data. The red × mark-

(A)

(B)

Fig. 9. GEMMs with intra-GPU collective speedup on (A) MI300X in CPX NPS4, and (B) MI355X in DPX NPS2 against GEMMs in SPX NPS1.

ers in Figure 8 indicate out-of-memory failures, which occurs only with naive sharding. Naive sharding exhausts capacity at batch sizes where both SPX/NPS1 and Starscream continue operating, a direct consequence of KV cache replication. d) Computation-Communication Overlap: As detailed in Section IV-F, we overlap computation with communication by customizing compute kernels to write outputs directly to the target NUMA region, eliminating separate collective kernels and their associated taxes. Figure 8C shows the performance benefit of this overlap over discrete kernels. This speedup compounds with Starscream’s other optimizations. e) Conservative Estimates: These results represent conservative lower bounds. We took the best-performing SPX/NPS1 kernel configuration and ported it to CPX/NPS4 with Starscream’s collectives and aggregation kernels, without tuning workgroup distribution, workgroup sizes, or flashattention segmentation for the CPX configuration. The occasional performance dips at certain batch sizes (e.g., 32, 64) reflect this lack of CPX-specific tuning. With dedicated optimization, we anticipate further speedups. C. GEMM with Intra-GPU Communication Figure 9 shows the performance of GEMM with intraGPU communication on MI300X (CPX/NPS4) and MI355X (DPX/NPS2) as speedup over the baseline SPX/NPS1. Each figure contains four quadrants: the top row shows GEMM followed by reduce-scatter; the bottom row shows all-gather followed by GEMM. The first column covers GEMM shapes occurring before the all-reduce in baseline tensor-parallel inference; the second column covers shapes occurring after the all-reduce. a) Connecting to Hierarchical Collectives: Recall from Section IV-D that Starscream decomposes all-reduce into intra-GPU reduce-scatter, inter-GPU peer-XCD all-reduce, and intra-GPU all-gather. In SPX mode, the computation is simply GEMM → inter-GPU all-reduce → GEMM. In CPX mode,

this becomes GEMM → hierarchical collective (reduce-scatter + all-reduce + all-gather) → GEMM. Since the inter-GPU all-reduce transfers the same amount of data in both cases, we isolate the additional CPX overhead by comparing: (1) GEMM-only in SPX against GEMM fused with intra-GPU reduce-scatter in CPX (top-left quadrant), and (2) GEMMonly in SPX against intra-GPU all-gather fused with GEMM in CPX (bottom-right quadrant). b) 2D Sharding via Collective Interchange: As discussed in Section IV-E, swapping GEMM+reduce-scatter with allgather+GEMM (and vice versa) enables 2D sharding with improved arithmetic intensity. The remaining two quadrants (topright, bottom-left) show this alternative. Comparing across quadrants, all four variants outperform SPX/NPS1, but the 2D-sharded configurations (swapped collectives) achieve additional speedup over their 1D counterparts—purely from better tile shapes and improved arithmetic intensity. c) Performance Across Batch Sizes: Figure 10 breaks down speedup by batch size for the all-gather+GEMM configuration on Llama3-405B FC2. Performance characteristics fall into three regimes: Low-latency (batch 1–16): In SPX/NPS1, all three matrices are distributed across 8 XCDs, and latency cannot be hidden. CPX/NPS4 leverages near-data compute with minimal remote access, achieving up to 1.41× speedup. • Throughput-optimized (batch 32–2048): Mid-range batches become compute-bound, allowing SPX to begin hiding memory latency. The performance gap narrows to approximately 1.08×, though CPX retains an advantage from superior bandwidth utilization. • Extreme-scale (batch >2048): Large batches exceed L2 capacity per XCD (FP16 on MI300X), causing sustained cache thrashing. Gains diminish at this scale; further improvement requires lower precision or larger L2 capacity. •

Fig. 10. Speedup of Llama3-405B in CPX NPS4 with AllGather followed by GEMM against SPX NPS1 as batch size increases.

VI. R ELATED W ORK a) Chiplet GPU Coherence: Prior work addresses coherence overhead in multi-chiplet GPUs assuming a shared coherence domain. HMG [47] extends coherence protocols hierarchically; CARVE [27] caches remote GPU data locally; CPElide [17] tracks inter-kernel dependencies to synchronize only when necessary, further improving performance. Starscream differs fundamentally: AMD’s CPX mode partitions a GPU into independent virtual GPUs with private L2 caches and disjoint memory, eliminating coherence overhead entirely. This introduces new challenges— reduced per-partition capacity and explicit communication— that Starscream addresses via orthogonal sharding. Even with zero coherence overhead, partitioning’s NUMA locality benefits (faster local HBM, 3× lower kernel launch latency) remain inaccessible to unified-memory designs like SPX NPS1. b) Parallelism Strategies: Distributing LLMs across GPUs relies on established techniques: tensor parallelism (TP) shards weight matrices with all-reduce reconciliation [51]; sequence parallelism (SP) shards activations along the sequence dimension for normalization and dropout [20], [28], [55]; context parallelism (CP) shards input sequences, exchanging KV tensors for full-context attention [1]; pipeline parallelism shards across layers [39]; expert parallelism routes tokens to sharded MoE experts [32]. Long-context LLMs have driven further refinements: Ring Attention [34] overlaps KV transfers with blockwise attention; Striped Attention [15] improves load balance via interleaved tokens; DeepSpeed-Ulysses [26] uses all-to-all for head-parallel attention; USP [21] and LoongTrain [22] unify these into hybrid schemes; DualPipe [19] overlaps forward and backward computation with communication to minimize pipeline bubbles. Helix [41] goes further, observing that a single uniform sharding is suboptimal: it disaggregates attention and FFN in time, reconfiguring the same GPUs between KV-parallel attention and TP FFN phases, bridged by a single all-to-all across a flat NVLink domain. These approaches share a common assumption: a homogeneous inter-GPU fabric. Starscream exploits a different axis of heterogeneity—the intra-GPU bandwidth hierarchy exposed by chiplet partitioning. Rather than replacing inter-GPU parallelism, Starscream composes in space: it retains the inter-GPU

layout unchanged and adds an orthogonal intra-die dimension across XCDs, exploiting the 10× bandwidth gap between onpackage USR and inter-die XGMI. Starscream further reinterprets parallelism through resource reduction: partitioning trades aggregate resources per virtual GPU for locality, and the optimal composition under this regime differs from the flat-topology optimum. This drives our design choices: where Helix and Ring Attention use uniform collectives, Starscream employs hierarchical collectives tuned to the NUMA topology and selects query- or KV-side movement per inference phase. The approaches are complementary: any inter-GPU strategy— Ring, Ulysses, or Helix—can serve as the outer layer that Starscream augments within each physical GPU. c) Partitioned and spatial accelerators: Beyond GPUs, spatial and wafer-scale accelerators expose fine-grained control over compute and data placement. Cerebras maps models onto a wafer-scale array of cores with software-managed onchip data movement [33], and a broad body of accelerator work explores tiling and dataflow mappings across spatially partitioned compute [16], [23]–[25], [29], [35], [36], [38], [43], [48], [49], [52], [54]. These systems demonstrate the value of locality-aware mapping, but they do so for bespoke architectures with programming models distinct from the GPU software stack. Starscream shows that comparable locality benefits are attainable within the commodity GPU programming model, by exploiting hardware partitioning modes rather than requiring a new class of accelerator. VII. C ONCLUSION Non-uniform memory access is an inherent reality in modern GPUs—whether from larger monolithic dies or chipletbased designs. AMD’s compute and memory partitioning modes expose this topology, providing knobs to exploit the underlying NUMA hierarchy. With perfect isolation, partitioning delivers its full potential. But when a model exceeds a single partition’s capacity, efficient mapping is essential to maximize reuse and minimize communication; without it, redundant parameters, KV caches, skinny matrices, and flat collectives erode the architectural gains. In this work, we presented Starscream, which reframes partitioning not as “more GPUs” but as a new level of the compute and memory hierarchy, by composing existing parallelism strategies hierarchically and applying each at the level of the topology where it is effective. Concretely, Starscream eliminates redundant memory usage, sustains arithmetic intensity through 2D GEMM sharding, structures collectives around the intra-die/inter-die bandwidth gap, and fuses computation and communication to further improve performance. On AMD Instinct MI300X and MI355X GPUs, these techniques deliver up to 4× faster attention and 1.43× faster GEMM execution, transforming GPU partitioning into a net advantage—using the same compute resources, with no additional hardware. As chiplet counts continue to grow and NUMA effects deepen, Starscream’s hierarchical sharding principles offer a blueprint for extracting performance from the topology of future partitioned accelerators.

R EFERENCES [1] “Context Parallel Package — Megatron Core — docs.nvidia.com,” https://docs.nvidia.com/megatron-core/developerguide/latest/user-guide/features/context parallel.html, [Accessed 20-072026]. [2] “PyTorch Symmetric Memory — PyTorch 2.13 documentation — docs.pytorch.org,” https://docs.pytorch.org/docs/2.13/symmetric memory.html, [Accessed 20-07-2026]. [3] S. Agarwal, B. Acun et al., “CHAI: Clustered head attention for efficient LLM inference,” in Proceedings of the 41st International Conference on Machine Learning (ICML), 2024. [4] J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai, “Gqa: Training generalized multi-query transformer models from multi-head checkpoints,” in Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, 2023, pp. 4895– 4901. [5] AMD, “Amd infinity architecture,” https://www.amd.com/en/ technologies/infinity-architecture, 2023, accessed: 2024. [6] AMD, “AMD Instinct MI300 Series: Compute and Memory Partitioning Modes,” Advanced Micro Devices, Inc., White Paper, 2024. [7] AMD, “AMD Instinct MI350 Series Accelerators,” https://www.amd. com/en/products/accelerators/instinct.html, 2024, accessed: 2024. [8] AMD, “CDNA Architecture: A New Era for Data Center and HPC Computing,” https://www.amd.com/en/technologies/cdna.html, 2024, accessed: 2024. [9] AMD, “CPX RCCL tuner,” https://rocm.blogs.amd.com/software-toolsoptimization/cpx-rccl-tuner/README.html, 2024, accessed: 2026-0728. [10] AMD, “Deep dive into the MI300 compute and memory partition modes,” https://rocm.blogs.amd.com/software-tools-optimization/ compute-memory-modes/README.html, 2024, accessed: 2026-07-28. [11] AMD, “MI300X GPU partitioning overview,” https://instinct.docs.amd. com/projects/amdgpu-docs/en/latest/gpu-partitioning/mi300x/overview. html, 2024, accessed: 2026-07-28. [12] AMD, “RCCL usage tips,” https://rocm.docs.amd.com/projects/rccl/en/ develop/how-to/rccl-usage-tips.html, 2024, accessed: 2026-07-28. [13] A. Arunkumar, E. Bolotin, B. Cho, U. Milic, E. Ebrahimi, O. Villa, A. Jaleel, C.-J. Wu, and D. Nellans, “Exploiting interposer technologies to disintegrate and reintegrate multicore processors,” IEEE Micro, vol. 36, no. 3, pp. 84–93, 2016. [14] A. Arunkumar, E. Bolotin, B. Cho, U. Milic, E. Ebrahimi, O. Villa, A. Jaleel, C.-J. Wu, and D. Nellans, “MCM-GPU: Multi-chip-module GPUs for continued performance scalability,” in Proceedings of the 44th Annual International Symposium on Computer Architecture (ISCA). New York, NY, USA: ACM, 2017, pp. 320–332. [15] W. Brandon, M. Mishra, A. Nrusimha, R. Panda, and J. R. Kelly, “Striped attention: Faster ring attention for causal transformers,” arXiv preprint arXiv:2311.09431, 2023. [16] J. Cai, Z. Wu, S. Peng, Y. Wei, Z. Tan, G. Shi, M. Gao, and K. Ma, “Gemini: Mapping and architecture co-exploration for large-scale DNN chiplet accelerators,” in 2024 IEEE International Symposium on HighPerformance Computer Architecture (HPCA), 2024, pp. 156–171. [17] P. Dalmia, R. S. Kumar, and M. D. Sinclair, “CPElide: Efficient multichiplet GPU implicit synchronization,” in Proceedings of the 57th IEEE/ACM International Symposium on Microarchitecture (MICRO), 2024. [18] DeepSeek-AI, “DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model,” arXiv preprint arXiv:2405.04434, 2024. [19] DeepSeek-AI, “DeepSeek-V3 technical report,” arXiv preprint arXiv:2412.19437, 2024. [20] J. Fang and S. Zhao, “Usp: A unified sequence parallelism approach for long context generative ai,” 2024. [Online]. Available: https://arxiv.org/abs/2405.07719 [21] J. Fang and S. Zhao, “USP: A unified sequence parallelism approach for long context generative AI,” arXiv preprint arXiv:2405.07719, 2024. [22] Y. Gu et al., “LoongTrain: Efficient training of long-sequence LLMs with head-context parallelism,” arXiv preprint arXiv:2406.18485, 2024. [23] S. Gudaparthi, S. Narayanan, R. Balasubramonian, E. Giacomin, H. Kambalasubramanyam, and P.-E. Gaillardon, “Wire-aware architecture and dataflow for CNN accelerators,” in Proceedings of the 52nd Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), 2019, pp. 1–13.

[24] S. Gudaparthi, S. Singh, S. Narayanan, R. Balasubramonian, and V. Sathe, “CANDLES: Channel-aware novel dataflow-microarchitecture co-design for low energy sparse neural network acceleration,” in 2022 IEEE International Symposium on High-Performance Computer Architecture (HPCA), 2022, pp. 876–891. [25] R. Hwang, T. Kim, Y. Kwon, and M. Rhu, “Centaur: A chiplet-based, hybrid sparse-dense accelerator for personalized recommendations,” in Proceedings of the ACM/IEEE 47th Annual International Symposium on Computer Architecture (ISCA), 2020, pp. 968–981. [26] S. A. Jacobs, M. Tanaka, C. Zhang, M. Zhang, S. L. Song, S. Rajbhandari, and Y. He, “DeepSpeed Ulysses: System optimizations for enabling training of extreme long sequence transformer models,” arXiv preprint arXiv:2309.14509, 2023. [27] A. Jaleel et al., “Combining HW/SW mechanisms to improve NUMA performance of multi-GPU systems,” in Proceedings of the 51st IEEE/ACM International Symposium on Microarchitecture (MICRO), 2018. [28] V. A. Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro, “Reducing activation recomputation in large transformer models,” in Proceedings of Machine Learning and Systems, D. Song, M. Carbin, and T. Chen, Eds., vol. 5. Curan, 2023, pp. 341–353. [Online]. Available: https://proceedings.mlsys.org/paper files/paper/ 2023/file/80083951326cf5b35e5100260d64ed81-Paper-mlsys2023.pdf [29] G. Krishnan, S. K. Mandal, M. Pannala, C. Chakrabarti, J.-S. Seo, U. Y. Ogras, and Y. Cao, “SIAM: Chiplet-based scalable in-memory acceleration with mesh for deep neural networks,” ACM Transactions on Embedded Computing Systems (TECS), vol. 20, no. 5s, pp. 1–24, 2021. [30] W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. Gonzalez, H. Zhang, and I. Stoica, “Efficient memory management for large language model serving with pagedattention,” in Proceedings of the 29th Symposium on Operating Systems Principles, 2023, pp. 611–626. [31] W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. E. Gonzalez, H. Zhang, and I. Stoica, “Efficient memory management for large language model serving with pagedattention,” in Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023. [32] D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen, “Gshard: Scaling giant models with conditional computation and automatic sharding,” arXiv preprint arXiv:2006.16668, 2020. [33] S. Lie, “Cerebras architecture deep dive: First look inside the hw/sw codesign for deep learning: Cerebras systems,” in 2022 IEEE Hot Chips 34 Symposium (HCS). IEEE, 2022, pp. 1–34. [34] H. Liu, M. Zaharia, and P. Abbeel, “Ring attention with blockwise transformers for near-infinite context,” in Proceedings of the 12th International Conference on Learning Representations (ICLR), 2024. [35] S. Luo, Y. Han, P. Li, J. Qin, J. Peng, Y. Zhao, Y. Cao, and T. Chen, “Mozart: Modularized and efficient MoE training on 3.5D wafer-scale chiplet architectures,” Advances in Neural Information Processing Systems, vol. 38, pp. 113 356–113 380, 2026. [36] M. Musavi, E. Irabor, A. Das, E. Alarcón, and S. Abadal, “Communication characterization of AI workloads for large-scale multi-chiplet accelerators,” in 2025 IEEE International Symposium on Circuits and Systems (ISCAS), 2025, pp. 1–5. [37] S. Naffziger, K. Viswanath, K. Wilcox, and N. Patki, “Chiplet-based high-performance computing: A survey,” IEEE Micro, vol. 43, no. 5, pp. 42–52, 2023. [38] P. S. Nalla, E. Haque, Y. Liu, S. S. Sapatnekar, J. Zhang, C. Chakrabarti, and Y. Cao, “CLAIRE: Composable chiplet libraries for AI inference,” in 2025 Design, Automation & Test in Europe Conference (DATE), 2025, pp. 1–7. [39] D. Narayanan, A. Harlap, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, P. B. Gibbons, and M. Zaharia, “Pipedream: generalized pipeline parallelism for dnn training,” in Proceedings of the 27th ACM Symposium on Operating Systems Principles, ser. SOSP ’19. New York, NY, USA: Association for Computing Machinery, 2019, p. 1–15. [Online]. Available: https://doi.org/10.1145/3341301.3359646 [40] NVIDIA, “TensorRT-LLM: High-Performance Large Language Model Inference,” https://github.com/NVIDIA/TensorRT-LLM, 2024, accessed: 2024. [41] NVIDIA, “Helix parallelism: Rethinking sharding strategies for

interactive multi-million-token LLM decoding,” arXiv preprint arXiv:2507.07120, 2025. [42] M. Osama, D. Merrill, C. Cecka, M. Garland, and J. D. Owens, “Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the gpu,” ser. PPoPP ’23. New York, NY, USA: Association for Computing Machinery, 2023, p. 429–431. [Online]. Available: https://doi.org/10.1145/3572848.3577479 [43] S. Pal, J. Liu, I. Alam, N. Cebry, H. Suhail, S. Bu, S. S. Iyer, S. Pamarti, R. Kumar, and P. Gupta, “Designing a 2048-chiplet, 14336core waferscale processor,” in 2021 58th ACM/IEEE Design Automation Conference (DAC), 2021, pp. 1183–1188. [44] D. Patterson and J. Hennessy, “The future of computing performance: Game over or next level?” Communications of the ACM, vol. 66, no. 4, pp. 62–70, 2023. [45] R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, A. Levskaya, J. Heek, K. Xiao, S. Agrawal, and J. Dean, “Efficiently scaling transformer inference,” https://arxiv.org/abs/2211.05102, 2022. [46] PyTorch Team, “Fully sharded data parallel: Pytorch 2.0 and beyond,” https://pytorch.org/docs/stable/fsdp.html, 2023, accessed: 2024. [47] X. Ren, D. Lustig, E. Bolotin, A. Jaleel, O. Villa, and D. Nellans, “HMG: Extending cache coherence protocols across modern hierarchical multi-GPU systems,” in Proceedings of the 26th IEEE International Symposium on High-Performance Computer Architecture (HPCA), 2020. [48] A. Shafiee, A. Nag, N. Muralimanohar, R. Balasubramonian, J. P. Strachan, M. Hu, R. S. Williams, and V. Srikumar, “ISAAC: A convolutional neural network accelerator with in-situ analog arithmetic in crossbars,” in Proceedings of the 43rd International Symposium on Computer Architecture (ISCA), 2016, pp. 14–26. [49] Y. S. Shao, J. Clemons, R. Venkatesan, B. Zimmer, M. Fojtik, N. Jiang, B. Keller, A. Klinefelter, N. Pinckney, P. Raina et al., “Simba: Scaling deep-learning inference with multi-chip-module-based architecture,” in Proceedings of the 52nd Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), 2019. [50] N. Shazeer, “Fast transformer decoding: One write-head is all you need,” arXiv preprint arXiv:1911.02150, 2019. [51] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro, “Megatron-lm: Training multi-billion parameter language models using model parallelism,” arXiv preprint arXiv:1909.08053, 2019. [52] Z. Tan, H. Cai, R. Dong, and K. Ma, “NN-Baton: DNN workload orchestration and chiplet granularity exploration for multichip accelerators,” in 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA), 2021, pp. 1013–1026. [53] O. A. Trifan, K. Sangaiah, M. Awad, M. Osama, S. Gudaparthi, A. Nicolau, A. Veidenbaum, and G. Dasika, “Eliminating multi-gpu performance taxes: A systems approach to efficient distributed llms,” 2025. [Online]. Available: https://arxiv.org/abs/2511.02168 [54] H. Wang, Z. Wang, H. Wang, J. Hou, T. Wei, C. Li, Y. Hu, and S. Yin, “WATOS: Efficient LLM training strategies and architecture coexploration for wafer-scale chip,” in 2026 IEEE International Symposium on High Performance Computer Architecture (HPCA), 2026, pp. 1–19. [55] B. Wu, S. Liu, Y. Zhong, P. Sun, X. Liu, and X. Jin, “Loongserve: Efficiently serving long-context large language models with elastic sequence parallelism.” New York, NY, USA: Association for Computing Machinery, 2024. [Online]. Available: https://doi.org/10. 1145/3694715.3695948 [56] Y. Zhao, A. Gu, R. Varma, L. Luo, C.-C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, A. Desmaison, C. Balioglu, P. Damania, B. Nguyen, G. Chauhan, Y. Hao, A. Mathews, and S. Li, “Pytorch fsdp: Experiences on scaling fully sharded data parallel,” in Proceedings of the VLDB Endowment, vol. 16, no. 12, 2023, pp. 3848–3860.

Related documents

• AMDKernelVault: Large-Scale Datasets and Agentic Training for AMD GPU Kernel Optimization#1069969
• Taming Bitwise Behavior in GPU Kernels with Tensor Core: Black-Box Reconstruction, Compiler Enforcement, and Static Verification#673469
• Divide and conquer: Scalable performance and energy in MCM GPUs#1187739
• Graph Traversal on Tensor Cores: A BFS Framework for Modern GPUs#259426
• GPUBreach: Privilege Escalation Attacks on GPUs using Rowhammer#1110228