CAIS：面向大模型张量并行的计算感知交换机内计算框架

本文介绍CAIS（Compute-Aware In-Switch computing）框架，旨在解决多GPU系统中张量并行的计算与通信隔离问题。该框架通过计算感知ISA扩展、合并感知的线程块协调、图级数据流优化器三大核心技术，实现了1.38倍的训练加速，为大规模AI基础设施提供新的设计范式。

背景：大模型训练中的通信瓶颈

随着大语言模型（LLM）规模扩张，单GPU无法满足需求，张量并行（TP）成为分布式训练核心策略，但频繁的集合通信操作成为性能瓶颈。传统NVLink SHARP（NVLS）技术通过交换机内计算加速通信，但以通信为中心的设计与LLM计算内核的内存语义存在根本性错配，导致计算与通信阶段隔离，资源利用率低下，限制重叠能力。

技术一：计算感知ISA与微架构扩展

CAIS定义计算感知指令集架构（ISA）并扩展交换机微架构。传统交换机仅处理转发，CAIS使交换机理解计算任务的内存访问模式（如读取、写入、原子操作），优化数据流动；微架构新增专用计算感知调度单元，根据GPU计算状态动态调整通信策略，确保数据到达时机与计算需求匹配。

技术二：合并感知的线程块协调

CAIS引入合并感知的线程块（TB）协调机制，分析各GPU TB执行进度，识别可合并的通信请求。当多个TB需访问相同/相邻数据时，协调到同一时间发起请求，充分利用交换机批处理能力；该机制动态调整，持续监控TB状态，预测通信需求并优化调度以最大化合并机会。

技术三：图级数据流优化器

CAIS的图级数据流优化器构建全局数据流视图，分析计算图的数据依赖关系，识别并行化机会。通过预取数据、延迟非关键通信、重新排序操作，实现跨内核的紧密重叠；该优化与张量并行特性契合，利用all-reduce等操作的数据局部性，提升流水线效率。

实验评估与性能表现

在主流LLM工作负载上，CAIS相比最先进的NVLS方案实现平均1.38倍端到端训练加速，与T3（无NVLS的计算-通信重叠方案）相比加速比达1.61倍。结果表明，需将计算与通信作为整体优化，CAIS消除了传统架构的计算-通信隔离，释放多GPU系统性能。

实际意义与未来展望

CAIS对大规模AI基础设施建设具有重要参考价值，展示了网络设备应成为计算生态系统积极参与者的设计范式。未来交换机可能集成更多计算能力，计算与网络融合将是下一代AI基础设施的重要趋势。结语：CAIS的1.38倍加速可降低训练成本或支持更大模型，为AI集群建设带来显著成本效益优势。