CAIS: Compute-Aware In-Switch Computing Framework for Tensor Parallelism in Large Models

This article introduces the CAIS (Compute-Aware In-Switch Computing) framework, which aims to solve the computation-communication isolation problem in tensor parallelism across multi-GPU systems. Through three core technologies—compute-aware ISA extension, merge-aware thread block coordination, and graph-level dataflow optimizer—the framework achieves a 1.38x training speedup, providing a new design paradigm for large-scale AI infrastructure.

Background: Communication Bottlenecks in Large Model Training

As the scale of large language models (LLMs) expands, a single GPU can no longer meet the demand. Tensor parallelism (TP) has become a core strategy for distributed training, but frequent collective communication operations have become a performance bottleneck. The traditional NVLink SHARP (NVLS) technology accelerates communication via in-switch computing, but its communication-centric design has a fundamental mismatch with the memory semantics of LLM computation kernels, leading to isolation between computation and communication phases, low resource utilization, and limited overlapping capability.

Technology 1: Compute-Aware ISA and Microarchitecture Extension

CAIS defines a compute-aware instruction set architecture (ISA) and extends the switch microarchitecture. Traditional switches only handle forwarding, while CAIS enables switches to understand the memory access patterns of computation tasks (e.g., read, write, atomic operations) to optimize data flow. The microarchitecture adds a dedicated compute-aware scheduling unit that dynamically adjusts communication strategies based on GPU computation status, ensuring that data arrival timing matches computation needs.

Technology 2: Merge-Aware Thread Block Coordination

CAIS introduces a merge-aware thread block (TB) coordination mechanism that analyzes the execution progress of each GPU TB and identifies mergeable communication requests. When multiple TBs need to access the same or adjacent data, they are coordinated to initiate requests at the same time, fully utilizing the switch's batch processing capability. This mechanism dynamically adjusts, continuously monitors TB status, predicts communication needs, and optimizes scheduling to maximize merging opportunities.

Technology 3: Graph-Level Dataflow Optimizer

CAIS's graph-level dataflow optimizer constructs a global dataflow view, analyzes the data dependencies of the computation graph, and identifies parallelization opportunities. Through prefetching data, delaying non-critical communication, and reordering operations, it achieves tight cross-kernel overlapping. This optimization aligns with the characteristics of tensor parallelism, leveraging data locality of operations like all-reduce to improve pipeline efficiency.

Experimental Evaluation and Performance

On mainstream LLM workloads, CAIS achieves an average 1.38x end-to-end training speedup compared to the state-of-the-art NVLS solution, and a 1.61x speedup compared to T3 (a computation-communication overlapping solution without NVLS). The results show that computation and communication need to be optimized as a whole—CAIS eliminates the computation-communication isolation of traditional architectures and unlocks the performance of multi-GPU systems.

Practical Significance and Future Outlook

CAIS has important reference value for large-scale AI infrastructure construction, demonstrating a design paradigm where network devices should be active participants in the computation ecosystem. Future switches may integrate more computing capabilities, and computation-network convergence will be an important trend for next-generation AI infrastructure. Conclusion: The 1.38x speedup from CAIS can reduce training costs or support larger models, bringing significant cost-effectiveness advantages to AI cluster construction.