Towards Compute-Aware In-Switch Computing for LLMs Tensor-Parallelism on Multi-GPU Systems

大规模大语言模型（LLM）推理与训练中的张量并行（Tensor Parallelism, TP）会频繁触发集体通信操作，这些操作已成为GPU间通信的主要瓶颈。以NVLink SHARP（NVLS）为代表的交换机内计算（in-switch computing）技术，虽可通过减少冗余数据传输来加速集体通信操作，但其以通信为中心的设计理念，却导致其通信模式与LLM计算核所要求的内存语义之间产生严重错配。这种错配使得计算与通信阶段彼此隔离，造成硬件资源利用率低下，且在多GPU系统中难以实现有效的计算-通信重叠。为克服这一局限，我们提出了CAIS——首个面向计算感知的交换机内计算框架，其核心思想是使交换机内的通信模式与计算核的内存语义需求保持一致。CAIS包含三项关键技术：（1）面向计算感知的指令集架构（ISA）与微架构扩展，以支持具备计算感知能力的交换机内计算；（2）面向请求合并优化的线程块（Thread Block, TB）协同机制，提升时序对齐精度，从而提高请求合并效率；（3）面向计算图层级的数据流优化器，实现跨计算核间的紧密重叠。在典型LLM负载上的实验评估表明：相较于当前最优、已集成NVLS的方案，CAIS平均端到端训练速度提升达1.38倍；相较于当前最优、支持计算-通信重叠但未利用NVLS的T3方案，CAIS提速达1.61倍。结果充分验证了CAIS在多GPU系统中加速张量并行的有效性。

在大规模LLM的张量并行（TP）训练与推理中，NVLink SHARP（NVLS）等片上交换计算技术虽能加速集体通信，但其纯通信导向的设计与LLM计算核对内存语义（如访存模式、数据重用、依赖关系）的需求存在根本性错配，导致计算与通信阶段割裂、资源利用率低、多GPU间计算-通信重叠不足——这是一个尚未被系统性解决的新问题。

首次提出‘计算感知的片上交换计算’（Compute-Aware In-Switch computing）范式：不将交换机视为黑盒通信加速器，而是将其微架构与指令集（ISA）扩展为可理解LLM计算内存语义（如TensorBlock级数据流、融合操作意图）的协同执行单元；通过计算感知ISA、合并感知线程块协调（merging-aware TB coordination）和图级数据流优化器三者协同，实现通信操作与计算内存在语义层面的对齐与紧耦合重叠。

在真实LLM训练负载（如Llama-2/3风格模型，TP=8）上评估，相较SOTA NVLS方案平均提速1.38×，相较不依赖NVLS但专注重叠的T3方案提速1.61×；实验覆盖训练吞吐、端到端时延、GPU利用率（via Nsight Compute/Systems）及通信-计算重叠率量化分析；论文未提开源代码，但硬件设计细节（如扩展ISA编码、TB调度协议）描述充分，为后续FPGA/ASIC原型实现与编译器支持（如MLIR集成）奠定基础；值得深入的方向包括：CAIS与流水线并行/序列并行的协同优化、面向MoE模型的动态稀疏通信适配、以及在CXL互连架构下的泛化。

1. T3: Tensor-Triggered Communication-Computation Overlap for Large Model Training (OSDI'23) 2. NVLink SHARP: Scalable Hierarchical Aggregation and Reduction Protocol (NVIDIA Whitepaper, 2021) 3. ColBERTv2: Effective and Efficient Retrieval via Lightweight Contextualization (ACL'22) ——虽非TP相关，但其向量级通信压缩思想启发CAIS的合并感知设计 4. PipeDream-2BW: Fast and Memory-Efficient Training of Large Models via Pipeline Parallelism (MLSys'24) 5. ZeRO-Offload: Democratizing Billion-Scale Model Training (MLSys'21)