TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives

TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives

Speakers: Size Zheng, Research Scientist, ByteDance; Jin Fang; Ningxin Zheng; Haibin Lin; Li-Wen Chang; Xin Liu

Conference: MLSys 2025

YouTube: https://slideslive.com/39043016

Overview

The rapid advancement and widespread adoption of large language models (LLMs) have driven an unprecedented demand for computational resources, particularly massive clusters of GPUs. Training and inference for state-of-the-art LLMs often necessitate tens of thousands of GPUs, making Model Flops Utilization (MFU) a critical metric for efficiency and cost-effectiveness. Even marginal improvements in MFU can translate into substantial cost reductions for organizations operating at this scale. A significant bottleneck in achieving high MFU is the overhead introduced by communication between GPUs, which, without optimization, can consume between 20% and 80% of the total execution time. This challenge is particularly acute in distributed training paradigms such as tensor parallel, pipeline parallel, data parallel, and sequence parallel.

This talk introduces TileLink, a novel distributed compiler designed to address the pervasive problem of communication overhead in large-scale ML systems. TileLink focuses on generating highly efficient compute-communication overlapping kernels, a technique crucial for hiding communication latency. The system proposes a unique approach that combines the benefits of existing overlapping strategies—decomposition-based and fusion-based—while mitigating their respective drawbacks. By introducing a set of tile-centric primitives, TileLink aims to simplify the development of complex, high-performance overlapping kernels, offering a significant reduction in implementation complexity without compromising performance.

TileLink represents a substantial step forward in distributed ML system optimization. Its ability to abstract away intricate hardware-specific details while exposing a simple, high-level programming interface empowers developers to craft sophisticated overlapping kernels with significantly less effort. The compiler's architecture, featuring a high-level Python AST transformation layer and a modified Triton compiler (Triton Distributed) at its core, demonstrates a powerful paradigm for bridging the gap between high-level programmability and low-level hardware efficiency. The presented results, showcasing substantial speedups across various LLM workloads on both NVIDIA and AMD GPUs, underscore TileLink's potential to dramatically improve the MFU and reduce operational costs for next-generation AI infrastructure.

Background

▶ Watch: LLM scale challenges and communication overhead problem (0:00)

The training and inference of large language models are inherently distributed processes, necessitating the coordinated effort of numerous computational devices. As models scale to billions and even trillions of parameters, they no longer fit onto a single GPU, leading to the widespread adoption of parallelism techniques such as tensor parallelism, pipeline parallelism, data parallelism, and sequence parallelism. While these strategies enable the execution of massive models, they introduce significant communication overheads. Data exchanges between GPUs—for example, in AllGather operations for tensor parallelism or ReduceScatter in data parallelism—can consume a dominant fraction of the total execution time, directly impacting the overall throughput and efficiency, measured by Model Flops Utilization (MFU). Reducing this communication latency is paramount for cost-effective large-scale ML.

Prior efforts to mitigate communication overhead primarily fall into two categories: decomposition-based and fusion-based approaches. Decomposition-based methods typically involve using multi-stream control to break down original operators into smaller sub-operators. These sub-operators can then be scheduled to run concurrently with communication operations if no data dependencies exist. This approach is generally straightforward to implement, often leveraging existing communication libraries like NCCL and compute kernels. However, its performance can be suboptimal because the overhead of managing multiple streams and the inherent granularity of decomposition may not fully hide communication latency. The discrete nature of these sub-operations often leaves "bubbles" in the execution timeline.

Conversely, fusion-based methods aim to achieve deeper optimization by fusing communication and computation kernels directly into a single, monolithic kernel. This tightly integrated approach allows for fine-grained control over data movement and computation scheduling, often leading to superior performance by effectively hiding communication latency within the fused kernel's execution. Examples include highly optimized, hand-tuned libraries that are specifically crafted for particular hardware and communication patterns. The significant drawback of fusion-based approaches, however, is their extreme complexity and development burden. Writing and maintaining such highly specialized, architecture-aware kernels requires deep expertise in low-level programming, GPU architecture, and distributed systems, making them prohibitively difficult for general application and rapid iteration. The challenge lies in creating a system that can deliver the performance benefits of fusion without the associated implementation complexity.

Key Findings

▶ Watch: TileLink's two-layer architecture and custom Triton (2:18)

TileLink's core contribution is a novel distributed compiler framework that combines the performance advantages of fusion-based compute-communication overlapping with the ease of implementation typically associated with decomposition-based methods. The system introduces a tile-centric primitive abstraction that allows developers to express complex overlapping patterns using significantly fewer lines of code compared to hand-tuned libraries. For instance, TileLink can achieve performance comparable to the highly optimized Flux library with only about one-third of the code complexity.

A key finding is the efficacy of TileLink's two-layer architecture. The high-level primitives provide a simple, semantic interface for managing data and control flow at the tile level, while the low-level primitives in the Triton Distributed compiler layer handle the intricate hardware-specific details. This separation of concerns, coupled with a Python-level Abstract Syntax Tree (AST) transformation for mapping high-level constructs to low-level implementations, is central to achieving both programmability and performance.

Experimentally, TileLink demonstrates significant performance gains across a wide range of LLM training and inference workloads on both NVIDIA and AMD GPUs. For fundamental operations like AllGather-GEMM and GEMM-ReduceScatter in dense models, TileLink not only matches but often surpasses the performance of PyTorch-based implementations (which leverage Cutlass/NCCL on NVIDIA and ROC-BLAS/RCCL on AMD). Crucially, it even shows speedups against highly optimized libraries like Flux, attributed to TileLink's superior swizzling techniques and intelligent choice between push and pull communication modes.

For more complex and dynamic workloads such as Mixture-of-Experts (MOE) layers and attention mechanisms, TileLink achieves even more dramatic improvements. It delivers "over one magnitude of speedup" compared to PyTorch for MOE layers by efficiently overlapping multiple operators like AllGather, GroupGEMM, ReduceScatter, TopK-Reduce, and local scatter. For attention mechanisms involving AllGather-KV, TileLink demonstrated "about two times speedup" over Ring Attention. Furthermore, its capabilities extend to inter-node communication, achieving significant speedups over PyTorch and comparable performance to Flux, even for challenging scenarios like very long context decoding (1 million context length) with excellent weak and strong scaling properties. The system also shows a "three times speedup to NCCL" for PCIe-based AllGather on L20 GPUs using NV-shmem primitives, and a remarkable "14 times speedup" for specific operators in DeepSeek-EP within an eight-node setup. These results collectively validate TileLink's design philosophy and its practical impact on distributed ML system efficiency.

Technical Deep Dive

▶ Watch: Primitive classification and AST translation mechanism (4:00)

TileLink's technical innovation lies in its hierarchical, primitive-based approach to generating compute-communication overlapping kernels. The architecture is composed of two primary layers: a high-level abstraction layer implemented via Python AST transformation and a low-level execution layer built on a modified Triton compiler, termed Triton Distributed.

At the high-level layer, TileLink exposes a small, intuitive set of nine primitives to the developer. These primitives are broadly categorized into:

1. Signal-related primitives: These manage control dependencies between producer and consumer tiles or peer relationships in distributed operations. Examples mentioned include consume_tile_wait and produce_tile_notify, which are crucial for synchronizing data readiness. peer_tile_wait is used for orchestrating ring-based communication patterns.
2. Data-related primitives: These govern the direction and mode of data transfer, offering choices between push mode and pull mode. In push mode, a rank proactively sends its data to remote ranks, typically involving a single synchronization barrier. In pull mode, a rank fetches data from remote ranks, often requiring two barriers: one to signal data readiness and another after transfer completion. The pull mode offers finer control over data arrival order, potentially enabling better overlapping, while push mode can be simpler but less flexible in scheduling.

The elegance of the high-level primitives lies in their simplicity; for instance, a wait primitive might only take a tile_ID as a parameter. However, translating this to actual hardware execution requires a wealth of underlying information. This is where the Python-level AST transformation comes into play. The core of this transformation is a sophisticated mapping process. The AST transformer analyzes the high-level primitive calls and, using static analysis or runtime lookup tables, retrieves the missing context from the tile_ID. This includes determining the rank_ID, barrier_ID, number of tiles per rank, and number of barriers per rank. For static workloads, simple affine transformations can map tile IDs to rank and barrier IDs. For dynamic workloads, such as Mixture-of-Experts (MOE) layers where dynamic routing causes varying relationships between ranks, TileLink allows lookup tables to be passed as parameters, enabling on-the-fly dependency calculation. For example, a high-level notify function might be transformed into a sequence of low-level primitives including synchronization, get_tile_id, and atomic operations.

The low-level layer is powered by Triton Distributed, a specialized version of the Triton compiler. This layer introduces approximately 36 low-level primitives that interact directly with the underlying hardware communication mechanisms. These primitives are divided into two groups:

1. Open-SHMEM primitives: These adhere to the Open-SHMEM standard, a specification for Partitioned Global Address Space (PGAS) programming. TileLink translates these primitives to vendor-specific implementations, such as NV-shmem for NVIDIA GPUs and ROC-shmem for AMD GPUs. This provides a portable abstraction over different shared memory architectures.
2. Non-Open-SHMEM primitives: These are designed to leverage specific hardware features not covered by Open-SHMEM, such as atomic operations and multi-memory features, allowing for fine-grained control and optimization tailored to particular GPU architectures.

Beyond the primitive design, TileLink incorporates advanced optimization techniques. Swizzling is a crucial optimization for cross-node operations, particularly exemplified in GEMM-ReduceScatter. In distributed scenarios, different ranks might process different data shards and initiate computation at varying times, leading to a "skew" in the global execution timeline. Swizzling primitives, such as produce_tile_swizzle and consume_tile_swizzle, are introduced to reorder data access and communication patterns. This ensures that communication aligns perfectly with computation, preventing bubbles (idle time) in the timeline and maximizing overlap. Developers can also integrate custom swizzling logic using basic transformations.

The design of TileLink's persistent kernels further illustrates its technical depth. For a GEMM-ReduceScatter kernel, the GEMM part is similar to a standard Triton GEMM, but with a single added produce_tile_notify primitive to signal data readiness. The subsequent ring-reduce kernel, which might utilize around 20 Streaming Multiprocessors (SMs), then employs consume_tile_wait and peer_tile_wait primitives to manage dependencies in a ring topology, ensuring efficient data exchange between neighboring ranks using a chosen data transfer mode (e.g., push mode). This modular and primitive-based composition allows for the construction of highly optimized kernels with minimal code changes, effectively bridging the gap between high-level declarative programming and low-level performance tuning.

Experimental Setup & Results

▶ Watch: Tile ID mapping and dynamic workload handling (5:00)

The evaluation of TileLink involved a comprehensive suite of experiments conducted on both NVIDIA GPUs (specifically L20 for some PCIe tests) and AMD GPUs, targeting a diverse set of 12 workloads representative of modern LLM training and inference scenarios. The primary baselines for comparison included PyTorch (which internally uses Cutlass and NCCL for NVIDIA, and ROC-BLAS and RCCL for AMD for compute and communication, respectively), Flux (a state-of-the-art communication library), vLLM, and Ring Attention.

For single-node performance, TileLink demonstrated strong capabilities on fundamental operations:

• AllGather-GEMM and GEMM-ReduceScatter, commonly used in dense models, showed good performance compared to PyTorch.
• Crucially, TileLink achieved noticeable speedups over Flux, a highly optimized library. This improvement was attributed to TileLink's advanced swizzling techniques and its strategic choice between push and pull communication modes, allowing for better compute-communication overlap.
• Similar speedups were observed on AMD GPUs for these same workloads, confirming the portability and effectiveness of TileLink across different hardware architectures.

For more complex and computationally intensive workloads, TileLink's benefits became even more pronounced:

• Mixture-of-Experts (MOE) layers, which involve intricate overlapping of multiple operators (AllGather, GroupGEMM, ReduceScatter, TopK-Reduce, local scatter), showed "over one magnitude of speedup" (i.e., >10x) compared to PyTorch. This highlights TileLink's ability to orchestrate complex dependencies efficiently.
• Compared to vLLM, TileLink also achieved speedups for MOE layers through its optimized overlapping.
• In attention layers, specifically for AllGather-KV, TileLink demonstrated "about two times speedup" when compared to Ring Attention, by effectively overlapping the AllGather operation with the attention computation itself.

Inter-node communication was another critical evaluation area:

• For AllGather and ReduceScatter operations, TileLink achieved "significant speedup to PyTorch" in cross-node scenarios. The talk noted that for these specific operations, communication performance is often not bounded by NIC bandwidth, allowing TileLink's optimizations to shine.
• Performance was found to be "comparable to Flux," positioning TileLink as a competitive solution even against highly specialized cross-node communication libraries.
• A particularly challenging experiment involved very long context decoding (batch size 1, 1 million context length) across eight nodes. TileLink exhibited "very good weak scaling" and "good strong scaling," effectively halving latency when the number of cards was doubled. This result is highly significant for the development of future reasoning models requiring extensive context.

Further specific results included:

• PCIe communication: On L20 GPUs for an AllGather example, TileLink, utilizing NV-shmem memory primitives and a low-latency protocol (different from NCCL's ring topology for PCIe), achieved "about three times speedup to NCCL."
• DeepSeek-EP operators: For specific operators within the DeepSeek-EP model, TileLink enabled the same functionality with "much less code" and delivered "about 14 times speedup" within an eight-node configuration.

The consistent and substantial speedups across diverse hardware, workloads, and communication patterns underscore TileLink's effectiveness in maximizing MFU by efficiently overlapping compute and communication, making it a powerful tool for large-scale ML deployments.

Practical Implications

▶ Watch: Key optimization tricks: Push/pull and swizzling (6:50)

TileLink offers profound practical implications for practitioners, infrastructure teams, model builders, and deployers operating at the cutting edge of large language model development. Its primary benefit is the significant improvement in Model Flops Utilization (MFU), which directly translates to substantial cost reductions for training and inference. Given that communication can consume 20-80% of total time in large distributed setups, TileLink's ability to effectively overlap compute and communication directly addresses one of the most significant bottlenecks in LLM scaling.

For model builders and researchers, TileLink drastically lowers the barrier to implementing high-performance distributed kernels. The high-level, tile-centric primitives abstract away the intricate details of low-level hardware interactions and complex synchronization logic. This means developers can focus more on algorithmic innovation and less on debugging arcane GPU assembly or communication protocols. The demonstrated reduction in lines of code (e.g., one-third compared to Flux) highlights a significant boost in developer productivity and maintainability for custom kernels. This agility is crucial in the fast-evolving landscape of LLM research, allowing for quicker experimentation with new architectures and distributed strategies.

Infrastructure teams stand to gain from TileLink's broad compatibility and optimization capabilities. By supporting both NVIDIA (NV-shmem) and AMD (ROC-shmem) GPUs, TileLink provides a more hardware-agnostic framework for distributed ML, potentially reducing vendor lock-in and allowing for more flexible hardware procurement. The compiler's ability to optimize inter-node communication, achieving performance comparable to state-of-the-art libraries like Flux, means that even clusters with high-latency network interconnects can see substantial throughput improvements. This enables more efficient utilization of existing hardware and potentially extends the lifespan of current infrastructure.

For ML deployers, the speedups demonstrated across various inference workloads (e.g., MOE, attention, long context decoding) directly translate to lower inference latency and higher throughput, which are critical for real-time applications and serving large user bases. The strong scaling capabilities for long context decoding are particularly relevant for emerging reasoning models that require processing vast amounts of information, paving the way for more powerful and responsive AI services.

However, there are also practical tradeoffs and limitations to consider. While TileLink simplifies kernel development, understanding the nuances of push versus pull communication modes and the importance of swizzling still requires a certain level of expertise in distributed systems to make optimal choices for specific workloads. The dynamic mapping for MOE layers, while powerful, introduces runtime overhead associated with lookup tables, which might need careful profiling for extremely latency-sensitive applications. Furthermore, while the system is open-source, its adoption will depend on community engagement and integration with popular ML frameworks. The current focus on Triton Distributed implies that users may need to adapt their existing PyTorch or TensorFlow code to TileLink's kernel definition format, which, while offering performance benefits, represents a shift from traditional framework-managed operations. Despite these considerations, TileLink provides a compelling solution to a critical problem, offering a clear path to more efficient and scalable LLM deployments.

Key Takeaways

• Significant MFU Improvement: TileLink drastically reduces communication overhead (20-80% of total time) in LLM training and inference by generating efficient compute-communication overlapping kernels, directly leading to substantial cost reductions.
• Simplified High-Performance Kernel Development: It combines the performance benefits of fusion-based approaches with the ease of implementation of decomposition-based methods, reducing code complexity by up to two-thirds compared to hand-tuned libraries like Flux.
• Two-Layer Architecture for Abstraction: The system features a high-level Python AST transformation layer with intuitive tile-centric primitives and a low-level Triton Distributed compiler layer that handles hardware-specific details (e.g., NV-shmem, ROC-shmem).
• Broad Performance Gains: TileLink delivers significant speedups across diverse LLM workloads, including dense models (AllGather-GEMM, GEMM-ReduceScatter), complex MOE layers (over 10x speedup to PyTorch), and attention mechanisms (2x speedup to Ring Attention) on both NVIDIA and AMD GPUs.
• Advanced Optimization Techniques: Key optimizations like dynamic mapping for MOE layers, strategic choice between push/pull communication modes, and swizzling for cross-node operations are crucial for maximizing overlap and minimizing bubbles in execution timelines.
• Strong Scaling for Future Models: Demonstrated excellent weak and strong scaling for very long context decoding (1 million context length across 8 nodes), paving the way for building more powerful and efficient reasoning models.

About the Speaker(s)

The talk was presented by Size Zheng from ByteDance. While specific titles for the co-authors (Jin Fang, Ningxin Zheng, Haibin Lin, Li-Wen Chang, Xin Liu) were not detailed in the transcript, Size Zheng's affiliation with ByteDance indicates a background in large-scale AI research and development, particularly in the context of advanced machine learning systems and infrastructure, given the focus of TileLink on optimizing LLM training and inference.