📝 Abstract

Large language models (LLMs) demand significant memory and computation resources. Wafer-scale chips (WSCs) provide high computation power and die-to-die (D2D) bandwidth but face a unique trade-off between on-chip memory and compute resources due to limited wafer area. Therefore, tensor parallelism strategies for wafer should leverage communication advantages while maintaining memory efficiency to maximize WSC performance. However, existing approaches fail to address these challenges. To address these challenges, we propose the tensor stream partition paradigm (TSPP), which reveals an opportunity to leverage WSCs’ abundant communication bandwidth to alleviate stringent on-chip memory constraints. However, the 2D mesh topology of WSCs lacks long-distance and flexible interconnects, leading to three challenges: 1) severe tail latency, 2) prohibitive D2D traffic contention, and 3) intractable search time for optimal design. We present TEMP, a framework for LLM training on WSCs that leverages topology-aware tensor-stream partition, traffic-conscious mapping, and dual-level wafer solving to overcome hardware constraints and parallelism challenges. These integrated approaches optimize memory efficiency and throughput, unlocking TSPP’s full potential on WSCs. Evaluations show TEMP achieves 1.7x average throughput improvement over state-of-the-art LLM training systems across various models.

💡 Analysis

Large language models (LLMs) demand significant memory and computation resources. Wafer-scale chips (WSCs) provide high computation power and die-to-die (D2D) bandwidth but face a unique trade-off between on-chip memory and compute resources due to limited wafer area. Therefore, tensor parallelism strategies for wafer should leverage communication advantages while maintaining memory efficiency to maximize WSC performance. However, existing approaches fail to address these challenges. To address these challenges, we propose the tensor stream partition paradigm (TSPP), which reveals an opportunity to leverage WSCs’ abundant communication bandwidth to alleviate stringent on-chip memory constraints. However, the 2D mesh topology of WSCs lacks long-distance and flexible interconnects, leading to three challenges: 1) severe tail latency, 2) prohibitive D2D traffic contention, and 3) intractable search time for optimal design. We present TEMP, a framework for LLM training on WSCs that leverages topology-aware tensor-stream partition, traffic-conscious mapping, and dual-level wafer solving to overcome hardware constraints and parallelism challenges. These integrated approaches optimize memory efficiency and throughput, unlocking TSPP’s full potential on WSCs. Evaluations show TEMP achieves 1.7x average throughput improvement over state-of-the-art LLM training systems across various models.

📄 Content

2026 IEEE International Symposium on High-Performance Computer Architecture (HPCA) TEMP: A Memory Efficient Physical-aware Tensor Partition-Mapping Framework on Wafer-scale Chips Huizheng Wang†∗, Taiquan Wei†∗, Zichuan Wang†, Dingcheng Jiang†, Qize Yang†, Jiaxin Liu†, Jingxiang Hou†, Chao Li‡, Jinyi Deng†B, Yang Hu†B, Shouyi Yin†§ †School of Integrated Circuits, BNRist, Tsinghua University, Beijing, China, 100084 ‡School of Computer Science and Engineering, Shanghai Jiao Tong University, Shanghai, China, 200240 §Shanghai Artificial Intelligence Laboratory, Shanghai, China, 200433 BCorresponding author, dengjinyi@mail.tsinghua.edu.cn; hu yang@tsinghua.edu.cn Abstract—Large language models (LLMs) demand significant memory and computation resources. Wafer-scale chips (WSCs) provide high computation power and die-to-die (D2D) bandwidth but face a unique trade-off between on-chip memory and compute resources due to limited wafer area. Therefore, tensor parallelism strategies for wafer should leverage communication advantages while maintaining memory efficiency to maximize WSC performance. However, existing approaches fail to address these challenges. To address these challenges, we propose the tensor stream partition paradigm (TSPP), which reveals an opportunity to leverage WSCs’ abundant communication bandwidth to alleviate stringent on-chip memory constraints. However, the 2D mesh topology of WSCs lacks long-distance and flexible interconnects, leading to three challenges: 1) severe tail latency, 2) prohibitive D2D traffic contention, and 3) intractable search time for optimal design. We present TEMP, a framework for LLM training on WSCs that leverages topology-aware tensor-stream partition, traffic- conscious mapping, and dual-level wafer solving to overcome hardware constraints and parallelism challenges. These inte- grated approaches optimize memory efficiency and throughput, unlocking TSPP’s full potential on WSCs. Evaluations show TEMP achieves 1.7× average throughput improvement over state-of-the-art LLM training systems across various models. I. INTRODUCTION Large language models (LLMs) have emerged as pivotal components in advancing artificial intelligence (AI) [6], [56], [73], [74], [92], [106], [140]. The LLM scaling law [141] highlights model size as a key performance driver, exem- plified by over 450× model size increase from GPT-2 [93] to DeepSeek [22]. Unfortunately, this rapid growth imposes significant demands on computation resources, making current monolithic devices increasingly difficult to provide sufficient transistor density for LLM training [32]. Wafer-scale chip (WSC) design, enabled by advanced pack- aging technologies like TSMC’s CoWoS [39], offers a promis- ing solution to mitigate these issues, by integrating numerous dies on a wafer-scale substrate. Compared to current board- level DGX systems, WSCs typically can offer 6× higher D2D bandwidth and 5× lower latency [36], [54], [109], benefiting from the finer and higher density interconnect pitches. *These authors contributed equally to this work. …… (a). Replication-relied TP ×

×

Die0 Die1 2.1 x mem usage in GPT-3 training Replication (b). Naive TSPP ×

×

Die0 Die1 2.7 x tail latency in GPT-3 training (c). Topology-aware TSSP Parallel-level

W0 W1

W0 W1 Logical-level Step 1 Advanced packaging Step 2 Step 3 Step 4 Performance Analysis Step 1 Step 2 Step 3 TEMP Potential Perf. Real Perf. WSC GPU DGX System TEMP Fig. 1. Illustrations for co-design features of the TEMP framework. However, different from current ASIC designs [23], [49], [77], [119], [121]–[124], chiplet [67], [98], [100], [107] and DGX-base designs [27], [115], [131], WSC-based systems face a unique architecture trade-off that both memory and compute resources are limited by the wafer’s area, typically 40,000mm2. In other words, increasing on-wafer memory capacity comes at the cost of compute resources. Therefore, LLM training on WSC-based systems necessitates the adoption of a memory-efficient tensor parallelism strategy. However, by re-examining existing tensor partition frame- works [9], [102], [143], [144], we identify that their paral- lelisms involve redundant tensor replication. As depicted in Fig.1 (a), although the weights are partitioned and stored separately across two dies, the activations remain replicated on both. To tackle the memory inefficiency issue, inspired by the distributed GEMM algorithms [14], [96], [117], we design a stream-style tensor partition mechanism, named TSPP. As shown in Fig. 1(b), TSPP partitions both input and weight ten- sors into non-overlapping sub-tensors, performs fine-grained sub-computations. While computing on a subset of these sub-tensors, the remaining sub-tensors are swiftly exchanged. TSPP offers two advantages: it eliminates tensor replication and enables the overlap of communication with computation. However, we identify that naively applying TSPP on WSCs is sub-optimal due to severe tail latency, as sho

This content is AI-processed based on ArXiv data.