MegaOps: A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

MegaOps: A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance Jinkun Lin

Byungsoo Oh

Rachee Singh

Cornell University Ithaca, NY, USA

Cornell University Ithaca, NY, USA

Cornell University Ithaca, NY, USA

arXiv:2611.16248v1 [cs.DC] 12 Nov 2026

Abstract

These tools trace kernel launches and analyze each kernel in isolation, allowing developers to attribute a performance anomaly or software bug to a specific kernel. Limitations of the kernel-per-operator paradigm. The kernelper-operator paradigm places a barrier between dependent kernels since the GPU runtime enforces dependencies at the granularity of entire kernels, either by executing kernels on the same CUDA stream in order or through GPU events between kernels on different streams. Consequently, a kernel cannot start until all thread blocks, or cooperative thread arrays (CTAs), of the kernels it depends on have finished, even when a part of its input is already available. This barrier causes three inefficiencies [2, 4]. First, when the last wave of a kernel’s CTAs does not fill every streaming multiprocessor (SM), the remaining SMs stay idle until this wave finishes. Second, a communication kernel cannot send any output of the preceding computation kernel until that kernel finishes, which serializes computation and communication. Finally, the next layer’s kernel cannot load its weights while the current layer’s kernel is running, which prevents the cross-kernel prefetching that weight-bound decode workloads rely on. Megakernels trade modularity for performance. Recent work eliminates kernel barriers by fusing the operators of an entire layer or an entire model into a single megakernel. MPK and ThunderKittens megakernels fuse the forward pass of an entire model into one kernel [4, 30], while FlashMoE fuses the expert computation and inter-GPU communication of a Mixture-of-Experts layer [2]. Inside a megakernel, a software scheduler assigns the tiles of all fused operators to SMs and manages the dependencies between individual tiles. This fine-grained scheduling lets the computation, communication, and weight loading of different operators overlap As a result, megakernels outperform the kernel-per-operator paradigm. However, fusing operators gives up modularity i.e., composability, interoperability, and debuggability, of the kernel-peroperator paradigm. Changing a single operator in a megakernel requires modifying the megakernel itself, which makes it costly to experiment with different kernel implementations or parallelization strategies. Reusing an existing kernel requires rewriting it in the format that the megakernel expects, and closed-source kernels from vendor libraries cannot be fused at all. Since the entire layer or model executes as one kernel, profiling and debugging tools that instrument kernel boundaries fail to attribute a performance anomaly or a failure to an operator. Must we choose between performance and modularity? In this work, we ask whether a kernel-per-operator ML framework can perform on par with megakernels. Kernel barriers are needed only because the GPU runtime enforces dependencies at the granularity of entire kernels. We posit that dependent kernels can instead

Modern ML frameworks execute models using the kernel-per-operator paradigm, which runs each operator as a separate GPU kernel. This paradigm is modular, since developers can compose various kernel implementations and profile each kernel in isolation. However, this paradigm enforces dependencies at the granularity of entire kernels, which places a barrier between dependent kernels. These barriers cause streaming multiprocessors to idle, serialize computation and communication, and prevent a kernel from prefetching its weights while the preceding kernel runs. Megakernels remove kernel barriers by fusing the operators of an entire layer or model into a single kernel, but they give up the modularity of the kernel-per-operator paradigm. In this work, we ask whether a kernel-per-operator ML framework can perform on par with megakernels. Our key insight is that removing kernel barriers does not require fusing kernels. We develop MegaOps, a framework that launches dependent kernels on separate CUDA streams and enforces each dependency on individual tiles of the intermediate tensor that the producer kernel writes and the consumer kernel reads. Since kernels synchronize on data tiles instead of on each other’s thread blocks, any two kernels that adopt the same data tiling can be composed without modifying either kernel. We evaluate MegaOps on dense Llama3 and Qwen3 models on single H100, A100, and L40 GPUs and on multi-GPU servers with 4×A100 and 8×H100. Our MegaOps framework matches the throughput of hand-tuned ThunderKittens megakernels, and achieves 1.09×–1.10× the decode throughput and 0.98×–1.10× the prefill throughput of the best-performing baseline, while reusing existing Triton, cuBLAS, and cuDNN kernels.

1

Introduction

Modern ML frameworks execute models using the kernel-per-operator paradigm [4]. In this paradigm, the framework represents an ML model as a computation graph of operators, e.g., matrix multiplication, attention, normalization, and collective communication, and runs each operator as a separate GPU kernel. GPU vendors and independent experts implement these kernels. For example, NVIDIA provides matrix multiplication kernels in cuBLAS, attention kernels in cuDNN, and collective communication kernels in NCCL [5, 20, 21], while experts have developed kernels like FlashAttention [6]. Modularity of the kernel-per-operator paradigm. Since each kernel interacts with the rest of the model only through its input and output tensors, the kernel-per-operator paradigm provides composability and interoperability which allow developers to pick and choose for each operator whichever kernel implementation best fits the requirements of their application. The kernel-per-operator paradigm also provides debuggability using tools like Nsight Systems and Nsight Compute [22, 23] which instrument kernel boundaries. 1

Jinkun Lin, Byungsoo Oh, and Rachee Singh

X

×

=

W

Y0

Y1

Y2

Y3

×

=

V

(a) Kernel-per-operator

SM0

Y0

SM1

Y1

SM2

Y2

Y3

Data tiles: rows of Y. Z0, Z1 need row 0; Z2, Z3 need row 1.

(b) Megakernel

(c) MegaOps

Z1

Z2

Z3

Z = YV Launch overhead

Y = XW

SMs idle

Stream

Data tile

Z0

Z0

Z3

Z1 Z2 time

Kernel barrier causes idle SMs and launch overhead

Y0

Kernel tile

operation. In contrast, a consumer in MegaOps waits on row 0 of 𝑌 regardless of how many producer CTAs computed it. This allows users to replace one producer kernel with a kernel that tiles 𝑌 differently, e.g., a Stream-K kernel [26], without modifying the consumer. Results. We implement MegaOps in Triton, reusing open-source kernels from TritonDistributed [39] and vLLM [14] and closedsource kernels from cuBLAS and cuDNN, and evaluate it on dense Llama and Qwen3 models with 0.6B to 70B parameters on single H100, A100, and L40 GPUs and on servers with four A100 and eight H100 GPUs. Our MegaOps framework achieves 0.89×–1.15× the throughput of ThunderKittens megakernels on the workloads they are hand-tuned for. Compared with the best-performing baseline in each setting among vLLM, TritonDistributed, and four megakernel systems, we achieve on average 1.10× (1.09×) the throughput in decode and 0.98× (1.10×) in prefill in single-GPU (multi-GPU) settings. Since kernels in MegaOps compose through data tiles, we can support different parallel schemes that share the same single-GPU compute code path, and require only 2.09× additional lines of code (LoC) relative to a PyTorch base implementation.

Y = XW Z = YV

Fused Y & Z Y0

Y3

Z2

Y0

Y3

Z2

Y1

Z0

Z3

Y1

Z0

Z3

Y2

Z1

Y2

Z1

time In-kernel tile-task dependencies

signal

wait

row 0 row 1 Data-tile semaphore of Y

Figure 1: Comparison of (a) kernel-per-operator, (b) megakernels and (c) MegaOps on back-to-back matmuls Y=XW, Z=YV. In (c) kernels are unfused and launched on independent streams to overlap. Y is partitioned into row data tiles; producer CTAs signal and consumer CTAs wait on semaphores before accessing each data tile. synchronize on individual tiles of the tensor that one kernel produces and the other consumes. This will enable a kernel to start as soon as the tiles it needs are ready, even while the kernel producing the remaining tiles is still running. Data-centric kernel overlap with MegaOps. Based on this insight, we develop MegaOps, a framework that runs each operator as a separate kernel and removes the barriers between dependent kernels. The MegaOps framework launches dependent kernels on separate CUDA streams such that the GPU runtime no longer orders the kernels. Instead, MegaOps enforces dependencies on individual tiles of the intermediate tensor that the producer kernel writes and the consumer kernel reads. To enable this synchronization, the user passes each kernel its semaphores and a function that maps the kernel’s own tiles to data tiles. Kernel developers adapt an existing kernel to MegaOps by adding a few lines of code that wait on or signal a semaphore at each tile access. Example. Figure 1 shows this design on two back-to-back matrix multiplications, 𝑌 = 𝑋𝑊 followed by 𝑍 = 𝑌𝑉 . In the kernel-peroperator paradigm (Figure 1(a)), tile 𝑌3 runs alone in the second wave of the first kernel, and the kernel computing 𝑍 cannot use the two idle SMs until 𝑌3 finishes. A megakernel fills these SMs by fusing both matrix multiplications (Figure 1(b)), whereas the MegaOps framework achieves the same schedule with two separate kernels (Figure 1(c)). The user partitions 𝑌 into two data tiles, one per row, and allocates a semaphore for each data tile. The CTAs computing 𝑌0 and 𝑌1 increment the semaphore of row 0 after writing their tiles, and the CTAs computing 𝑍 0 and 𝑍 1 start on the idle SMs as soon as this semaphore indicates that row 0 is complete, while 𝑌3 is still running. Data-centric vs. CTA-centric synchronization. Synchronizing on data tiles is the key idea that allows MegaOps to compose operator kernels. CuSync and FlashOverlap also overlap dependent kernels on separate streams, but they synchronize on the CTAs of the producer kernel instead of the data it produces [9, 12]. In Figure 1, a CuSync consumer computing 𝑍 0 must know which producer CTAs computed row 0 of 𝑌 . FlashOverlap keeps track of tiles in the order the producer computes them and supports only a matrix multiplication followed by a collective communication

2

MegaOps Framework

The MegaOps framework consists of two components, a data-tile kernel interface that lets dependent kernels synchronize on tiles of the tensors between them (§2.1, §2.2), and a runtime called AsyncStream that launches kernels on multiple CUDA streams so that the GPU runtime does not order them (§2.3). By the interface of a kernel, we mean the arguments through which a caller launches it. In the traditional kernel-per-operator paradigm, this interface consists of the kernel’s input and output tensors and a few configuration parameters. This interface does not reveal how the kernel divides its work among CTAs. MegaOps extends this interface with arguments that tell a kernel when each part of its input is ready and where to record that each part of its output has been written. We use the back-to-back matrix multiplications in Figure 1 as a running example.

2.1

A Data-centric Principle for Kernel Composition

In the kernel-per-operator paradigm, the interface of a kernel consists of its input and output tensors. Since a caller interacts with a kernel only through these tensors, it does not depend on how the kernel is implemented or tiled. As a result, one kernel can replace another without changes elsewhere in the model. Overlapping dependent kernels, as proposed by MegaOps, appears to break this composability. When a producer and a consumer kernel run concurrently on separate streams (Figure 1(c)), a consumer CTA must wait until the producer CTAs that write its input have finished. If the consumer waits on producer CTAs, it must know how the producer assigns tiles to CTAs, which couples the tilings of the two kernels. MegaOps avoids this coupling by synchronizing on data instead of CTAs. Data tile and readiness semaphore. MegaOps asks the user to partition the tensor that connects the two kernels into data tiles, independent of the tiling of either kernel. The consumer waits 2

MegaOps : A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

def kernel(..., in_semaphore: Tensor # data tile semaphore 3 # for input tensor 4 in_map: Callable # kernel-tile-->data-tile 5 # N-D coordinate mapping 6 in_target: int # target semaphore value 7 # for readiness 8 in_sema_extent: tuple # N-D data-tile 9 # coordinate extent 10 # similar arguments for output 11 out_semaphore: Tensor 12 out_map: Callable 13 out_sema_extent: tuple 14 # optional tile size configurations 15 tile_config: dict 16 # optional remap to reschedule CTAs 17 # in a different order 18 cta_remap: Callable): ... Listing 1: MegaOps kernel interface

(0, )

1 2

64

(0, 0) (0, 1)

(0, 0) (0, 1) (0, 2)

64

(1, 0) (1, 1) (1, 2)

(1, 0) (1, 1) (1, 2)

64

64

64

64

64

𝐵

64

64

64

𝐶

Figure 2: Kernel tiles and coordinates for 𝐶 = 𝐴𝐵 with tile sizes of 64 on row, column, and reduction dimensions. The blue kernel tiles are accessed when computing 𝐶 tile (0, 1). until data tiles it reads are ready. Each data tile is assigned a multidimensional coordinate. The user creates a semaphore for each data tile to track if the producer has fully written it. Inspired by Event Tensor [13], semaphores are ordinary integer tensors owned and managed by users, so that semaphore hangs can be debugged on the user side by inspecting their values with a timeout mechanism we provide. Users then adapt both the producer and consumer kernels to the data tiling via our kernel interface, so that a producer CTA updates the semaphore to signal that a data tile is ready, and a consumer CTA checks the semaphore before it reads the data tile. As a result, dependent kernels compose freely as long as they are adapted to the same data tiling.

2.2

(0, 1)

(0, 2)

128

(1, 0)

(1, 1)

(1, 2)

128

128

128

(1, )

(0, 0)

Kernel tile

(2, )

(0, )

Data tile

(3, )

to expose their tiling scheme and assign each kernel tile a multidimensional coordinate in natural language as part of the function documentation, so that the user knows which part of the data each kernel tile accesses. For example, Fig. 2 shows the kernel tiling and coordinates in a matmul kernel with 64 as tile sizes in all dimensions. The kernel developer can also choose to expose coarser-grained kernel tiles than the actual kernel implementation to reduce synchronization overhead, like merging tiles in the same row into a bigger tile. Using this documentation, the user passes mappings and extents that translate kernel tiles into data tiles. The mapping functions, in_map and out_map, map a kernel-tile coordinate to the coordinate of the first data tile it intersects, and the extents, in_sema_extent and out_sema_extent, specify how many data tiles the kernel tile spans in each dimension. Together, they identify the semaphore entries that a CTA synchronizes on. Our kernel-side contract is that whenever the kernel accesses a kernel tile, it uses the mapping and extent to find the corresponding data tiles, and then waits on in_semaphore until each entry reaches in_target, or signals out_semaphore. Here, in_target specifies the number of kernel tiles needed to fully produce the data tile, which can be ≥ 1 as a data tile may likewise have multiple intersecting kernel tiles. Finally, the user can pass a tile_config to customize the kernel’s tile sizes, akin to the block sizes in many Triton kernels. Optionally, kernels can expose a cta_remap interface that remaps CTA IDs to let users change the CTA processing order, e.g., to follow the order in which the producer makes data tiles ready (§3.1). Example. Listing 2 shows how our interface can be used to implement a back-to-back matmul 𝑌 = 𝑋𝑊 , 𝑍 = 𝑌𝑉 . First, the user decides the data tiling scheme and tile size of the connecting intermediate tensor 𝑌 , which uses row-wise tiling with size 64. Next, to run the first matmul, the user tunes the tile size with tune_config and gets the kernel tile sizes. For simplicity, the example assumes the tuned tile size can divide or be divided by 64. To adapt the producer to 𝑌 ’s data tiling, the user maps the kernel tile to 𝑌 ’s data tile, where the kernel has tile size ktile_m and 𝑌 has dtile_m. ktile_coord[0] * ktile_m // dtile_m maps the coordinate of the first element, which is passed to out_map. To complete the mapping, the user computes the number of data tiles a kernel tile covers with cdiv(ktile_m, dtile_m) and passes it to out_sema_extent. The user computes the number of signals every data tile will receive and stores it in y_target for the consumer. Fig. 3 shows the mapping. The user adapts the consumer kernel similarly on the input side.

(0, 0) (0, 1) (0, 2)

(1, 0) (1, 1)

𝐴

(0, 0)

Figure 3: Mapping kernel tiles of 𝑌 = 𝑋𝑊 to row-wise data tiles in Listing 2. Each kernel tile intersects two data tiles, e.g., kernel tile (1, 1) maps to data tiles 2 and 3, encoded by out_map([1,1]) = 1 · 128 // 64 = 2 with extent of 2. Each data tile receives y_target= 3 signals.

=

×

128

Kernel Interface

To synchronize on data tiles, a kernel needs to know how the user has partitioned each tensor into data tiles. The kernel must wait on or signal the corresponding semaphores before it reads or after it writes a data tile. Our kernel interface, summarized in Listing 1, lets users pass this information to the kernel. We currently support synchronization on one input tensor and one output tensor, so a kernel accepts a single in_semaphore and a single out_semaphore, which contain the data-tile semaphores of these tensors. Kernels typically partition the output tensor into tiles and assign each tile to a CTA, so each CTA accesses a tile of every input and output tensor, which we call a kernel tile. We require kernels 3

Jinkun Lin, Byungsoo Oh, and Rachee Singh

1

@triton.jit def matmul_kernel(A, B, C, M, N, K, 3 TILE_M, TILE_N, TILE_K, sync_ctx): 4 _kernel_begin_cta(sync_ctx) #+ 5 pid = tl.program_id(0) 6 nn = tl.cdiv(N, TILE_N) 7 _pids = (pid // nn, pid % 8 _pids = _cta_remap(sync_ctx, _pids) #+ 9 pid_m, pid_n = _pids #+ 10 _ktile_in = (pid_m,) #+ expose row-tiling 11 _ktile_out = (pid_m, pid_n)#+ grid-tiling 12 _wait_input_cta(sync_ctx, _ktile_in) #+ 13 acc = tl.zeros((TILE_M, TILE_N), ...) 14 for k in range(0, K, TILE_K): 15 a = tl.load(...) 16 b = tl.load(...) 17 acc = tl.dot(a, b, acc) 18 tl.store(..., acc) 19 _signal_output_cta(sync_ctx, _ktile_out)#+ Listing 3: Kernel-side matmul integration example. Starting from an ordinary Triton GeMM (tile addressing omitted), the developer adds the sync_ctx argument and the lines marked #+: L4 signals the start of the kernel; L8 reorders CTA IDs; L12 waits with 𝐴’s kernel tile coordinate; L19 signals with 𝐶’s kernel tile coordinate.

m, hidden_size = x.shape

1 2

2

# pick rowwise plain data tiling scheme dtile_m = 64 # data tile size 5 y_sema = zeros(cdiv(m, dtile_m), 6 dtype=torch.int32, device=x.device) 3 4

7

# adapt producer t_cfg = tune_config('matmul', x, W) 10 ktile_m = t_cfg['tile_m'] # kernel tile size 11 y = matmul(x, W, t_cfg, 12 out_semaphore=y_sema, 13 out_map=f'lambda ktile_coord:ktile_coord[0] 14 * {ktile_m} // {dtile_m}', 15 out_sema_extent=cdiv(ktile_m, dtile_m)) 16 # number of ktiles per dtile 17 y_target = cdiv(y.shape[1], t_cfg['tile_n']) 18 * cdiv(dtile_m, ktile_m) 8 9

19

# adapt consumer t_cfg = tune_config('matmul', y, V) 22 ktile_m = t_cfg['tile_m'] 23 z = matmul(y, V, t_cfg, 24 in_semaphore=y_sema, 25 in_map=f'lambda ktile_coord:ktile_coord[0] 26 * {ktile_m} // {dtile_m}', 27 in_sema_extent=cdiv(ktile_m, dtile_m)), 28 in_target=y_target) Listing 2: Back-to-back matmul of 𝑌 = 𝑋𝑊 , 𝑍 = 𝑌𝑉 in MegaOps. cdiv is ceiling division. We assume either data tile size divides the kernel tile sizes or the reverse for simplicity. We discuss indivisible cases later. 20 21

our functions support any arbitrary Triton functions, so users may even precompute these function values and store them in a tensor and load them back when needed. This is useful for highly irregular tiling where no closed-form formula exists for extents or targets. Utility functions to simplify user programming in plain tilings. In a plain tiling, a tensor is partitioned regularly into consecutive chunks of a fixed size in each dimension. Irregular tilings are rare. Most often a plain tiling scheme is used for both data and kernel tiles. The user programming effort can be further reduced in this case using our utility functions. We observe that the programming pattern is homogeneously matching two tile grids with different grid sizes on the same tensor. For example, in Fig. 3 computing 𝑌 ’s kernel-to-data tile extent can be abstracted as computing how many dashed grid cells a solid grid cell maps to. We therefore provide the following utility functions using this grid abstraction so that users do not need to calculate the formula manually. • calc_extent(𝑑𝑡𝑖𝑙𝑒, 𝑠ℎ𝑎𝑝𝑒, 𝑘𝑡𝑖𝑙𝑒) returns a lambda function (in a Python string) that accepts a kernel tile coordinate and calculates the extent, i.e., the number of data tiles that intersect with it, potentially dynamically. 𝑑𝑡𝑖𝑙𝑒 and 𝑘𝑡𝑖𝑙𝑒 represent the data and kernel tile sizes, respectively, and 𝑠ℎ𝑎𝑝𝑒 represents the full shape of the tensor. • calc_target(𝑑𝑡𝑖𝑙𝑒, 𝑠ℎ𝑎𝑝𝑒, 𝑘𝑡𝑖𝑙𝑒) similarly returns a lambda function that accepts a data tile coordinate and returns the number of kernel tiles that intersect. Using them, the calculation formulas for y_target and extents in Listing 2 can both be replaced with a simple function call. Supporting GPUDirect peer-to-peer (P2P) memory access. Communication kernels, especially in the scale-up domain, often

Kernel integration. Integrating a kernel to support our interface is straightforward and follows programming patterns common in many of today’s kernel design conventions such as prologue, epilogue, and tile scheduler from CUTLASS [18]. We provide utility functions to facilitate the integration, including _wait_input_cta for input and _signal_output_cta for output semaphore synchronization, and the kernel only needs to invoke them with the kernel tile coordinates. The utility functions convert them to data tile ones and perform the corresponding synchronization under the hood. Our prototype is based on Triton, and Listing 3 shows the integration of a GeMM kernel as an example. It only requires adding a few lines of code, demonstrating that the effort is small to integrate a kernel in MegaOps. Handling irregular and indivisible tile mapping. We further extend our interface to accept a callable for in/out_sema_extent and in_target, which accept the kernel-tile and data-tile coordinate, respectively. This allows the value to vary from tile to tile by computing it dynamically, similar to the in/out_maps. This is useful when different data tiles have different numbers of kernel tiles intersecting, which happens often with indivisible tile sizes. A function-based extent and target computation can address the issue. To increase flexibility, the user can add any additional arguments at the kernel launch and pass them to these functions. Note that 4

MegaOps : A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

rely on direct P2P memory access to reduce local data copy and achieve efficient communication. The kernel typically relies on symmetric memory [24, 27], where all ranks allocate a tensor of the same size and can address peer tensors directly. To support this pattern, we prepend a rank dimension to our data and kernel tile coordinates to identify remote tiles. The semaphores for a symmetric memory tensor will also be stored using symmetric memory to let a kernel synchronize the readiness of a remote data tile. Supporting closed-source kernels. Triton kernels may not always have the best performance, and other kernels with higher performance may not have source code available for the above integration, such as the cuBLAS and copy engine kernels. Inspired by prior decomposition techniques [34, 35, 38], we split the output tensor into a few chunks and launch one smaller kernel, which we call a subkernel, for each chunk. We conceptually regard these subkernels as a coarse-grained set of CTAs of the undecomposed kernel, and invoke the provided wait and signal functions on the kernel tiles they access at subkernel boundaries. To make this happen, we run subkernels in two separate streams in a round-robin order to avoid inter-subkernel barriers, and wrap each with a tiny wait and signal synchronization kernel to control when to start a subkernel and to notify completion. We launch a proxy kernel on a third stream to interact with these synchronization kernels, so that we can invoke _wait_input_cta and _signal_output_cta in the proxy before starting a subkernel and after one finishes.

Green Contexts [17], which can assign streams to a GPU partition, and launch kernels to corresponding streams based on kernel types. Automatic memory management. PyTorch’s CUDA memory pool assumes a tensor is used on a single stream, so under multistreaming it may reuse a tensor whose kernel is still running. Its stream-recording workaround incurs host-device synchronization and is incompatible with CUDA graphs. Furthermore, PyTorch’s memory pool is nondeterministic in the reuse outcome, which causes issues when reusing symmetric memory. This is because each symmetric memory tensor has its paired peer tensors, and all ranks must make the same decision and reuse the same paired ones. We address both automatically with delayed tensor release and a consistent symmetric memory pool. We cap the number of concurrent kernels at 𝑁 using CUDA events and hold a reference to each tensor a kernel uses until 𝑁 more kernels have launched, by which time that kernel must have finished. For symmetric tensors, we additionally hold the reference until all peers are known to have passed the 𝑁 -kernel mark. We rely on user annotation of cross-peer data dependency to trace peer progress without extra synchronization (details in Appendix A). Finally, to guarantee peer consistency, symmetric tensors are allocated from a dedicated pool that reuses only exact-size matches and returns the earliest freed tensor, so that all ranks make the same decision. We also run Python’s garbage collection (GC) manually in the same schedule across peers for consistency, and freeze prior generations to reduce the overhead. Appendix A details the design.

2.3

3

AsyncStream Runtime

We develop a runtime called AsyncStream to overlap the kernels by dispatching to separate streams in a round-robin schedule. The number of streams is configurable, trading between interference and overlap. Exposed to the outside, AsyncStream provides an abstraction as if all MegaOps kernels are run under one single stream, except that the stream imposes no inter-kernel barrier. AsyncStream then offers common wait and block stream synchronization primitives, so that it can interact with other regular streams and coexist with existing unintegrated kernels, maximizing compatibility and enabling gradual integration. In doing so, AsyncStream addresses the following challenges to achieve correct and efficient scheduling and overlap. Deadlock-free scheduling. It guarantees that a kernel launched earlier can secure hardware resources before later ones to avoid deadlocks, using the same wait kernel technique from CuSync: a wait kernel is launched on the same stream before every kernel launch to delay the current kernel from being scheduled until after the prior ones have secured their hardware resources. Empirically, we observe that CUDA schedules all CTAs of a kernel before scheduling the next. Thus, we emit a signal from the prior kernel whenever any CTA has started using _kernel_begin_cta (Listing 3), which the wait kernel waits on. Overlapping communication and computation kernels with spatial partitioning. Without spatial partitioning, the earlier kernels, especially computation ones, will often occupy the whole GPU, leaving no resources available to overlap the second kernel. We address this problem by spatially carving out a small number of SMs for running communication kernels, as a few SMs are sufficient to supply sufficient bandwidth for overlapping [28, 31]. We use

Overlap-aware Optimization

Launching dependent kernels on separate streams in MegaOps removes kernel barriers, but by itself it does not deliver the performance of megakernels. We find two causes of this behavior. First, existing kernels are written and tuned under the assumption that they own the GPU, so they do not take advantage of the opportunities that overlap creates. For example, a kernel could prefetch its weights while the preceding kernel is still running, and it could process its tiles in the order in which the producer finishes them, but existing kernels do neither (§3.1). Second, overlapping kernels compete for SMs, and the CUDA scheduler decides how to place their CTAs. A poor placement leaves SMs idle or slows down the execution (§3.2). We address the first cause by reordering instructions and the second with a mechanism that regulates CTA placement. Since the best settings for both depend on how the kernels in a model overlap, we tune them jointly with a graph-level autotuner (§3.3).

3.1

Dependency-aware Reordering and Prefetching

The existing kernel-per-operator paradigm assumes all inputs are ready before kernel launches. Kernels thus often adopt a swizzled CTA schedule to optimize data locality. Under kernel overlap, this assumption breaks: some CTAs that are scheduled early may stall on unavailable inputs while others with input ready are scheduled to start later, creating pipeline bubbles. Inspired by CuSync, we expose a cta_remap interface (Listing 1&3) that allows users to align CTA scheduling with producer data ready order. For example, in an allgather–GeMM overlap, the push-based allgather [35] produces 5

Jinkun Lin, Byungsoo Oh, and Rachee Singh

4

local shards (via local memcpy) before remote ones. Local GeMM CTAs can run first by remapping 𝑘𝑡𝑖𝑙𝑒 [0] to (𝑘𝑡𝑖𝑙𝑒 [0]+𝑑)%𝑚, where 𝑑 is the first local tile index and 𝑚 the number of row tiles. In lowbatch decoding, GeMM weight loading dominates execution time. Kernel overlap allows us to prefetch weights during the preceding kernel’s execution. We lift a tunable number of weight tile loads to the beginning of the kernel code, before semaphore waits, to prefetch. More prefetches hide more latency but require earlier kernel launches (so more streams) and consume more SMEM/registers. Our graph pass autotunes the prefetch depth.

5 3.2

Evaluation

We answer the following questions in the evaluation. • Can MegaOps match megakernel performance in model serving, in both single- and multi-GPU settings? • Can MegaOps match handwritten libraries in computationcommunication overlap workloads? • How much is the coding effort on the user side? • Is graph tuning essential to obtain high performance?

Overlap Regulation with SMEM Adjustment

Overlapping kernels interfere with each other in CTA-to-SM scheduling. We observe that spatially sharing the same SM between two overlapping kernels does not always happen, even when the SM has enough resources left. This causes severe performance loss to kernels with a smaller number of waves, such as persistent GeMMs or under small batch sizes: they cannot use the SMs occupied by prior overlapping kernels, and CUDA eagerly schedules all CTAs to the remaining SMs at the moment. When prior kernels finish, the released SMs become idle, wasting resources. We address it with CTA launch blocking. We observe that prior kernels are often at the end of execution when the next kernel starts launching CTAs. Therefore, we block some launches of the next kernel’s CTAs until prior kernels finish, using an SMEM padding technique to force a CTA admission control. In detail, as each SM has a limited SMEM capacity, we can increase SMEM of a kernel to limit the number of colocated CTAs. By allocating a large enough SMEM, the unoccupied SM subset no longer has enough SMEM to host all CTAs, which forces CUDA to stop launching new CTAs. We design an occupancy limit parameter 𝑐, the maximum number 𝐶 + 1 where 𝐶 of CTAs allowed, and set the SMEM allocation to 𝑐+1 is the GPU SMEM capacity. Our graph pass autotunes 𝑐. On the other hand, when two kernels use complementary hardware resources, such as tensor cores versus CUDA cores, it is beneficial to reduce SMEM usage of a kernel to encourage spatial SM sharing. To do so, we tune the number of stages, a common knob existing in modern kernels, to reduce SMEM usage, and autotune it in our graph pass.

3.3

Implementation

We prototype our design using Triton and TritonDistributed. The functions in the interface are implemented as Triton functions, taking advantage of Triton’s support for passing Triton functions as kernel arguments. Lambda functions are supported by translating them to a Triton function at runtime. PyTorch symmetric memory is used to support peer memory and multicast address mapping. Appendix B lists the kernels we support and details our scheduling optimization of closed-source kernels. We are compatible with and rely on CUDA graphs to reduce launch overhead.

5.1

Evaluation Setup

We use both single-GPU and single-node multi-GPU settings to evaluate our system. Single-GPU settings consist of A100 (40GB SXM), L40, and H100 (80GB HBM3), while multi-GPU settings consist of 4×A100 GPUs fully connected via direct P2P NVLink and 8×H100 connected via NVLS, with 100 GB/s and 450 GB/s unidirectional bandwidth, respectively. We describe the multi-GPU results on 8×H100 in this section and defer the 4×A100 results to Appendix C.5, which show similar trends. We use CUDA 13.1 and PyTorch 2.11.0 for all experiments. Workloads. We evaluate the performance of a single batch of requests serving (§5.2) using long and short contexts, over a set of batch sizes. The long decode uses 1024 for prompt length and 1024 for output length, while the short decode uses 128 for prompt length and 64 for output length. We use Qwen3-0.6B, 1.7B, 8B, Llama-3.2-1B, and Llama-3.1-70B for this evaluation, which cover a variety of dense models. To evaluate the effect of computationcommunication overlap more directly, we use allgather-GeMM overlap from a TritonDistributed benchmark (§5.3) which runs MLP allgather and up-projection overlap under tensor parallelization. The weight shapes come from Llama-3.1-8B, 70B, and 405B, and Qwen2-72B, GPT3-175B with varying batch sizes, covering both small and large shapes. Baselines. We compare against vLLM, TritonDistributed, megakernel compilers and frameworks MPK, MegaTritonKernel [38, 42] (referred to as MegaTD) from TritonDistributed, and VDCores [8], and specialized megakernels ThunderKittens including both the singleGPU low-latency and the 8-GPU high-throughput versions. We also run a no-overlap version of our system by disabling multi-streaming and removing the prefetch and SMEM padding. For clarity, we omit systems on prefill or decode plots when they fail to run or perform poorly (Appendix C.1). For communication-computation overlap, we compare against AsyncTP [35] and TritonDistributed. We run each system 5× and report the median normalized against vLLM. We report the speedup against the best other systems on the bar plots.

Graph-level Autotuning

The parameters mentioned above are sensitive to how kernels overlap and hence cannot be determined on a per-kernel basis. We thus develop a graph-level tuning algorithm that runs all kernels to optimize these parameters. To search efficiently, we adopt an iterative tuning approach where we vary only one parameter and fix all others at a time. Each parameter enumerates over a set of predefined candidate values and profiles performance to select the best one. To reduce overhead for transformer models that have homogeneous layers, only the first few layers are used for profiling. We also tune the number of streams and number of decomposed chunks, as well as other customized knobs required to be tuned in an overlap environment. 6

MegaOps : A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

1.40×

1.36×

1.36×

1.32×

1.32×

1.5 0.89×

1.0

1.0

0.5

0.5

0.0

MegaTD

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

1.09×

1.05×

1.05×

1.18×

BS=1 BS=2 BS=4 BS=8 BS=16

VDCores

1.5 1.04×

ThunderKittens

Qwen3-1.7B 1.04×

1.00×

1.01×

Qwen3-8B 1.20×

1.0 0.5 0.0

0.98×

BS=1 BS=2 BS=4 BS=8 BS=16

0.98×

0.97×

0.96×

1.09×

1.0

Not supported

1.5

MPK

Llama-3.2-1B

Not supported

Ours (no overlap)

Not supported

H100

Normalized throughput (vs. vLLM)

Ours

Qwen3-0.6B

0.5 0.0

BS=1 BS=2 BS=4 BS=8 BS=16

Figure 4: Short decode throughput on H100 using 128 prompt tokens and 64 output tokens. Numbers above bars indicate the speedup of “Ours” over the best of the other systems. A100 and L40 results are in Fig. 9 in Appendix. MegaOps implementation and parallel schemes. We implement the dense Llama and Qwen3 model families using our layerlevel interface. For multi-GPU execution, we implement two parallel schemes: an allreduce (AR) scheme from the traditional Tensor Parallelism [29] and an allgather-reducescatter (AG-RS) scheme from ThunderKittens megakernels [31]. AR is used for decode, which has a smaller batch. For prefill, we manually profile both schemes to pick a batch size threshold: AG-RS is used from BS ≥ 8 for Qwen38B and BS ≥ 2 for Llama-3.1-70B on 8×H100, and from BS ≥ 2 on 4×A100. Most intermediate tensors use a row-wise data tiling with a row tile size of 128 (details in Appendix C.1).

8xH100

Normalized throughput (vs. vLLM)

Ours (no overlap)

MPK

MegaTD

Qwen3-8B 1.18×

1.20×

1.17×

1.16×

1.10×

1.09×

1.0

1.0

0.5

0.5

0.0

ThunderKittens

Llama-3.1-70B

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

1.06×

1.07×

1.10×

1.01×

BS=1 BS=2 BS=4 BS=8 BS=16

8xH100

Figure 5: Long decode throughput. “BS” denotes global batch size. MPK does not support Llama-3.1-70B.

End-to-end Model Offline Prefill and Decode

Single-GPU result. Fig. 4 shows the short-context decode-only throughput normalized against vLLM. Compared to the highly specialized megakernel baselines, the MegaOps framework achieves comparable performance, within 0.96×–1.09× of VDCores and 0.89× of ThunderKittens. For long-context decode, we outperform ThunderKittens by 1.33×, while VDCores crashes due to KV cache index integer overflow (Fig. 11 in Appendix). Note that ThunderKittens is a low-latency megakernel that is compatible only with Llama3.2-1B at BS = 1, and VDCores uses specialized schedules for each model at BS ≤ 8. The MegaOps framework achieves a comparable throughput without manual tuning (§3.3) on the models and batch sizes they are hand-tuned for, demonstrating that our generality only comes at a minimal cost of performance. Appendix C.8 analyzes where the MegaOps framework gains and loses against VDCores and the non-overlapped version with a per-kernel trace breakdown analysis. Compared to vLLM, a non-megakernel system, we achieve speedup consistently across all models and batch sizes, and on average obtain 1.27× (1.14×–1.48×) higher short-context decode throughput than vLLM on H100. The trend holds on A100 and L40 (Appendix C.2), with average speedups of 1.23× and 1.15×, respectively. As vLLM uses different kernels than us, to separate out this effect and measure the impact of overlap alone, we also run a version of MegaOps with overlap turned off. The results show that overlap contributes significantly to the speedup: without it, the MegaOps framework reaches 1.05×, 1.12×, and 1.07× of vLLM’s throughput, and overlap adds an additional 1.21×, 1.10×, and 1.07×, on H100, A100, and L40 GPUs, respectively. We achieve smaller speedups as model sizes grow because the kernel body execution time increases and dominates the overall latency. Appendix C provides more detailed analysis. In prefill settings, we only achieve comparable throughputs with baselines for a similar reason.

Normalized throughput (vs. vLLM)

5.2

Ours

Qwen3-8B 1.5 1.0

1.04×

1.12×

1.13×

1.22×

Llama-3.1-70B 1.36× 1.01×

0.98×

1.06×

1.10×

0.5

0.5 0.0

1.0

0.94×

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

BS=1 BS=2 BS=4 BS=8 BS=16

Figure 6: Prefill throughput. “Ours” uses AG-RS for BS ≥8 for Qwen3-8B and BS ≥2 for Llama-3.1-70B, while AR is used for the rest of the BS. Legend is the same as Fig. 5. Multi-GPU decode. Fig. 5 shows the long-context decode throughput on 8×H100. On average, our MegaOps framework outperforms vLLM by 1.11× (1.01×–1.20×). The improvement is slightly lower than in the single-GPU decode, as computation in decode has only 1 wave with no room for the additional all-reduce operator to be hidden with overlap. The megakernel compilers all achieve consistently lower throughput than vLLM in the batch sizes we evaluate, while ThunderKittens performs poorly as it is designed for larger batch sizes rather than our small-batch settings. Multi-GPU prefill. Unlike the single-GPU case, multi-GPU prefill benefits from overlap, as there are more computation kernel waves to hide communication. Fig. 6 shows that we outperform vLLM by 1.18× (1.01×–1.48×) on 8×H100. Compared to the non-overlapped version, we achieve 1.24× higher throughput. We find that the best parallel scheme depends on the batch size for each model. At BS=1, AG-RS reaches only 0.72×–0.90× of AR’s throughput, while at BS=16 it is 1.08×–1.13× faster than AR. This demonstrates the importance of flexible programmability that allows for easy switching between parallel schemes depending on the workload. Compared to the 8-GPU ThunderKittens megakernel on Llama-3.1-70B, our MegaOps framework achieves 0.94×–1.15× of its throughput. Our major gain comes from a faster QKV projection and RoPE computation and better overlap between allgather and gate- and up-projection than ThunderKittens’s broadcast-based allgather implementation, which demonstrates the importance of individual kernel tuning. 7

Jinkun Lin, Byungsoo Oh, and Rachee Singh

5.3

Table 1: Model forward pass lines of code breakdown. The ratio row divides each column by the Base lines.

Computation-communication Overlapping

Fig. 7 shows that our green context design can overlap computation and communication efficiently. It achieves comparable performance to the best baselines in larger shapes where communication can be fully hidden. In smaller shapes, we achieve up to 18% higher throughput because AsyncTP and TritonDistributed both use copy engine, which has a fixed exposed cost relatively high in small messages, while serial PyTorch uses NCCL for allgather, which does not fully utilize the GPU SMs and results in longer allgather time. Our framework automatically switches to kernel-based allgather for small shapes, and runs it on the computation green context partition to fully utilize GPU SMs. This demonstrates the importance of switching kernels adaptively according to workloads. Our composability makes the switching seamless, which enables autotuning that can decide the best combination automatically.

5.4

In Out Alloc Tune Remap Other Total Lines Ratio

Usability Evaluation

We quantify the user-side effort by counting the lines of code of our model implementations. We categorize each line of code into two main categories: base code and MegaOps addition. The latter is further divided into six subcategories. Below are the classification rules: • Base: code irrelevant to MegaOps, including kernel invocations with their arguments, buffer preparation such as tensor allocation, reshapes and initialization, backend dispatch, etc. • In/Out: code preparing and passing the input/output semaphore arguments (in_semaphore, out_semaphore, map, target, and extents). • Alloc: code constructing semaphore tensors, including choosing data tile sizes and sizing semaphore tensors. • Tune: code invoking the autotuner and passing the resolved tile_config to the launch. • Remap: CTA remapping (§3.1). • Other: stream synchronization, symmetric-memory handling, profiling, and tile_config bookkeeping. Each line is assigned to one category, and a line mixing base and MegaOps content is classified to MegaOps categories conservatively.1 Table 1 shows that our framework adds 2.1× lines per line of the original base code, with a total of 1767 additional lines on top of Base’s 847 lines, confirming that only a small amount of code is needed to use the MegaOps interface. This code size can be further reduced by providing more helpers for common patterns, similar to calc_target and calc_extent (§2.2). For example, 0.29× lines of this are for autotuning, which contains repeated argument lists as the kernel calls; 0.20× are for allocation, which contains similar size calculation patterns. The remaining is mostly for constructing the input and output mappings and can also be simplified with similar utility functions as those for target and extent. We leave these improvements to future work.

5.5

MegaOps addition

Base

Graph-level Tuning Ablation

To evaluate the overlap-aware optimizations (§3) and whether graph-level tuning is necessary to make them effective, we rerun the 1 Blank lines, comments, function signatures, assertions, and debug code are excluded. The large line count is due to our linter, which breaks down long lines into one argument per line.

8

847 490 611 1.00 0.58 0.72

171 0.20

244 0.29

111 0.13

140 0.17

1767 2.09

short-context decode on H100 with graph tuning entirely skipped, and with one optimization removed from the graph tuning search at a time while still tuning the rest. Fig. 8 shows the slowdown of the per-step decode time relative to the full graph tuning. Without graph tuning, the MegaOps framework uses kernel configurations autotuned on each kernel in isolation, and the decode step becomes more than 10% slower in certain cases, confirming the effectiveness of graph-level tuning. We next ablate the effect of the proposed knobs in the graph-level tuning. Without tuning prefetching, the decode step slows down by 2-6% for small models, where GeMMs are small and the prefetching covers a larger fraction of the total weight tiles. Disabling both SMEM padding and stage tuning slows down the decode step by up to 6%. Appendix C.6 further shows that disabling CTA remapping slows down the multi-GPU prefill by up to 1.32×.

5.6

Composability Case Study

As mentioned in §2.2, as long as two kernels are adapted to the same data tiling scheme, they can compose freely. To demonstrate this concretely, we illustrate with how two different MLP TP strategies, AG-RS and allreduce (AR), can share the same code path of the MLP computation (pseudocode in Listing 4 in Appendix). The AG-RS path runs AG→MLPCompute→RS, while the allreduce (AR) path runs NoOP→MLPCompute→AR. By all using the row-wise data tiling scheme for the intermediate tensors, namely the input and output of AG, MLPCompute, RS, and AR, the same MLPCompute function can be shared between the two communication paths. This allows us to easily switch between AG-RS and AR without changing other parts of the code path, demonstrating the composability of our system. In fact, we add AR after the AG-RS path is implemented, and adding AR for both MLP and Attention blocks only required us to edit in total 346 Python LoC in the user program, including weight sharding logic, control-flow branching logic, and safety checks.

6

Related Work

Several lines of work study how to reduce the overhead in the traditional kernel-per-operator scheme. Megakernels pack individual kernels together to reduce overhead. Some works develop highly specialized megakernels for certain workloads or hardware, such as ThunderKittens megakernels, FlashMoE, Perseus, and UniEP [2, 25, 30, 31, 33, 41]. To improve portability, several works design general frameworks to generate megakernels. MPK, Event Tensor, MegaTritonKernels, and TileRT can compile a user program resembling PyTorch to a megakernel [1, 4, 13, 40, 42]. VDCores [8] proposes a fine-grained megakernel programming model where users schedule the compute or memory operations per SM. However, they all fuse

8xH100

Normalized throughput (vs. Serial PyTorch)

MegaOps : A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

Ours

LLaMA-3.1-8B 1.5

0.92× 1.18×

1.13×

0.88×

1.5

1.0

1.0

0.5

0.5

0.0

1K

TritonDistributed

LLaMA-3.1-70B 0.99×

2K 4K 8K M (tokens)

16K

0.0

1.08×

0.93×

0.92×

0.93×

AsyncTP

Qwen2-72B 1.01×

1.5 1.0

0.90× 1.02×

1.02×

GPT-3-175B 0.97×

0.95×

1.0

2K 4K 8K M (tokens)

16K

0.0

0.97×

0.99×

1.00×

LLaMA-3.1-405B 1.00×

1K

2K 4K 8K M (tokens)

16K

0.0

1.05×

1.05×

1.06×

1.00×

1.01×

1K

2K 4K 8K M (tokens)

16K

1.0 0.5

0.5

0.5 1K

1.03×

1K

2K 4K 8K M (tokens)

16K

0.0

Figure 7: Allgather-GeMM overlap.

Slowdown (%)

Full graph tune (min-max) No prefetch

Llama-3.2-1B

Qwen3-0.6B 10

5

5

5

2.5 0

0 1

2 4 8 Batch size

16

1

2 4 8 Batch size

16

1

2 4 8 Batch size

Conclusion

ML frameworks face a choice between the modularity of the kernelper-operator paradigm and the performance of megakernels. We observed that megakernels gain their performance by removing kernel barriers, which does not require fusing kernels. MegaOps runs each operator as a separate kernel, launches dependent kernels on separate CUDA streams, and synchronizes them on tiles of the tensors that connect them. Since kernels synchronize on data instead of on each other’s CTAs, users can compose and reuse kernels freely and debug each operator in isolation. MegaOps matches the throughput of hand-tuned ThunderKittens megakernels on their target workloads, showing that ML frameworks can reach megakernel performance without fusing operators and without giving up the modularity of the kernel-per-operator paradigm.

Qwen3-8B

10

0

7

No SMEM padding or stage tuning Pre graph tune

16

Figure 8: Effect of graph-level tuning. “Pre graph tune” does not run graph-level tuning. The slowdown is relative to full graph tuning, with min and max plotted as error bars.

Acknowledgement

into one monolithic kernel and lose interoperability and debuggability. Our MegaOps framework retains these properties while achieving comparable performance to megakernels. In multi-GPU settings, computation-communication overlap is another approach to reducing the kernel-per-operator overhead. Most of them design specific fused kernels or overlap schemes for certain operator combinations, such as AG+GeMM, GeMM+RS, GeMM+AR, Dispatch+GroupGeMM, and GroupGeMM+Combine [3, 9, 10, 15, 16, 34–37]. CoCoNet, TritonDistributed, Mercury, Syncopate, and ParallelKittens design domain-specific languages (DSLs) for generating customized kernel overlap or fusion [7, 11, 28, 32, 39]. They all focus on co-designing computation and communication to achieve high performance but often introduce kernel coupling. We instead focus on a general interface for overlapping any number of both computation and communication kernels with minimal coupling. Our interface is closest to CuSync and Event Tensor, which also let users express tile-level dependencies [12, 13]. CuSync builds dependencies on CTA IDs, so a consumer kernel must know which producer CTAs produce the tile it loads. This couples the tilings of the two kernels. Event Tensor represents a dependency as an event that counts the completion of a set of producer tasks. Users map the task coordinates of both kernels onto a shared event tensor. Since an event records that tasks have finished rather than that a region of a tensor is ready, the event layout is a convention that each pair of kernels must agree on. MegaOps instead ties every semaphore to a data tile of the tensor itself, so the data tiling alone defines the contract between any two kernels. Event Tensor also fuses its device functions into a single persistent kernel, so closed-source kernels cannot participate and per-kernel profiling tools no longer apply, whereas MegaOps launches each operator as a separate kernel and overlaps them at runtime.

We thank Riley Bakes for implementing a multi-threaded version of semaphore signal and wait functions and Chi-Sheng Liu for an early exploration of the SM scheduling issue.

References [1] Tilert: Tile-based runtime for ultra-low-latency llm inference. https://github. com/tile-ai/tilert. [2] Osayamen Aimuyo, Byungsoo Oh, and Rachee Singh. FlashMoE: Fast Distributed MoE in a Single Kernel. Advances in Neural Information Processing Systems, 38:100676–100699, 2025. [3] Li-Wen Chang, Wenlei Bao, Qi Hou, Chengquan Jiang, Ningxin Zheng, Yinmin Zhong, Xuanrun Zhang, Zuquan Song, Chengji Yao, Ziheng Jiang, et al. Flux: Fast software-based communication overlap on gpus through kernel fusion. arXiv preprint arXiv:2406.06858, 2024. [4] Xinhao Cheng, Zhihao Zhang, Yu Zhou, Jianan Ji, Jinchen Jiang, Zepeng Zhao, Ziruo Xiao, Zihao Ye, Yingyi Huang, Ruihang Lai, et al. { MPK } : A compiler and runtime for { Mega-Kernelizing } tensor programs. In 20th USENIX Symposium on Operating Systems Design and Implementation (OSDI 26), pages 1909–1926, 2026. [5] Sharan Chetlur, Cliff Woolley, Philippe Vandermersch, Jonathan Cohen, John Tran, Bryan Catanzaro, and Evan Shelhamer. cuDNN: Efficient primitives for deep learning. arXiv preprint arXiv:1410.0759, 2014. [6] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022. [7] Yue Guan, Xinwei Qiang, Zaifeng Pan, Daniels Johnson, Yuanwei Fang, Keren Zhou, Yuke Wang, Wanlu Li, Yufei Ding, and Adnan Aziz. Mercury: Unlocking multi-gpu operator optimization for llms via remote memory scheduling. In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles, pages 1046–1061, 2025. [8] Zijian He, Adrian Sampson, Yiying Zhang, and Zhiyuan Guo. VDCores: Resource decoupled programming and execution for asynchronous GPU. arXiv preprint arXiv:2605.03190, 2026. [9] Ke Hong, Xiuhong Li, Minxu Liu, Qiuli Mao, Tianqi Wu, Zixiao Huang, Lufang Chen, Zhong Wang, Yichong Zhang, Zhenhua Zhu, et al. Efficient and adaptable overlapping for computation and communication via signaling and reordering. In Proceedings of the 21st European Conference on Computer Systems, pages 1894– 1911, 2026. [10] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: 9

Jinkun Lin, Byungsoo Oh, and Rachee Singh

[35] Yifu Wang, Horace He, Less Wright, Luca Wehrstedt, Tianyu Liu, and Wanchao Liang. [distributed w/ TorchTitan] introducing async tensor parallelism in PyTorch. PyTorch Developer Forums, https://discuss.pytorch.org/t/distributed-wtorchtitan-introducing-async-tensor-parallelism-in-pytorch/209487, September 2024. [36] Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, Li-Wen Chang, et al. Comet: Fine-grained computation-communication overlapping for mixture-of-experts. Proceedings of Machine Learning and Systems, 7, 2025. [37] Chenggang Zhao, Shangyan Zhou, Liyue Zhang, Chengqi Deng, Zhean Xu, Yuxuan Liu, Kuai Yu, Jiashi Li, and Liang Zhao. Deepep: an efficient expertparallel communication library. https://github.com/deepseek-ai/DeepEP, 2025. [38] Size Zheng. Mega triton kernel. https://github.com/ByteDance-Seed/Tritondistributed/blob/main/docs/getting-started/megakernel/megakernel.md, 2025. [39] Size Zheng, Wenlei Bao, Qi Hou, Xuegui Zheng, Jin Fang, Chenhui Huang, Tianqi Li, Haojie Duanmu, Renze Chen, Ruifan Xu, et al. Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler. arXiv preprint arXiv:2504.19442, 2025. [40] Size Zheng, Jin Fang, Xuegui Zheng, Qi Hou, Wenlei Bao, Ningxin Zheng, Ziheng Jiang, Dongyang Wang, Jianxi Ye, Haibin Lin, et al. Tilelink: Generating efficient compute-communication overlapping kernels using tile-centric primitives. Proceedings of Machine Learning and Systems, 7, 2025. [41] Size Zheng, Xuegui Zheng, Li-Wen Chang, and Jidong Zhai. UniEP: Unified expert-parallel MegaKernel MoE for LLM training. In Proceedings of the 35th International Symposium on High-Performance Parallel and Distributed Computing, HPDC ’26, pages 387–401, New York, NY, USA, 2026. ACM. [42] Size Zheng, Xuegui Zheng, Hanshi Sun, Qi Hou, Wenlei Bao, Shiyu Li, Haojie Duanmu, Jin Fang, Chenli Xue, Chenhui Huang, et al. Ditron: Distributed multilevel tiling compiler for parallel tensor programs. arXiv preprint arXiv:2605.02953, 2026.

System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509, 2023. [11] Abhinav Jangda, Jun Huang, Guodong Liu, Amir Hossein Nodehi Sabet, Saeed Maleki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Saarikivi. Breaking the computation and communication abstraction barrier in distributed machine learning workloads. In Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, pages 402–416, 2022. [12] Abhinav Jangda, Saeed Maleki, Maryam Mehri Dehnavi, Madan Musuvathi, and Olli Saarikivi. A framework for fine-grained synchronization of dependent gpu kernels. In 2024 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), pages 93–105. IEEE, 2024. [13] Hongyi Jin, Bohan Hou, Guanjie Wang, Ruihang Lai, Jinqi Chen, Zihao Ye, Yaxing Cai, Yixin Dong, Xinhao Cheng, Zhihao Zhang, et al. Event tensor: A unified abstraction for compiling dynamic megakernel. Proceedings of Machine Learning and Systems, 8:1917–1933, 2026. [14] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with PagedAttention. In Proceedings of the 29th Symposium on Operating Systems Principles (SOSP), 2023. [15] Nandor Licker, Kevin Hu, Vladimir Zaytsev, and Lequn Chen. fabric-lib: Rdma point-to-point communication for llm systems. Proceedings of Machine Learning and Systems, 8:169–185, 2026. [16] Hao Liu, Matei Zaharia, and Pieter Abbeel. Ringattention with blockwise transformers for near-infinite context. In International Conference on Learning Representations, volume 2024, pages 3992–4008, 2024. [17] NVIDIA. 4.6. green contexts — cuda programming guide. https://docs.nvidia. com/cuda/cuda-programming-guide/04-special-topics/green-contexts.html. [18] NVIDIA. Cutlass. https://github.com/nvidia/cutlass. [19] NVIDIA. Stream memory operations. https://docs.nvidia.com/cuda/cuda-driverapi/cuda_driver_api/group__CUDA__MEMOP.html. [20] NVIDIA. cuBLAS: Basic linear algebra on NVIDIA GPUs. https://docs.nvidia. com/cuda/cublas/, 2026. [21] NVIDIA. NVIDIA collective communications library (NCCL). https://developer. nvidia.com/nccl, 2026. [22] NVIDIA. NVIDIA Nsight Compute. https://developer.nvidia.com/nsightcompute, 2026. [23] NVIDIA. NVIDIA Nsight Systems. https://developer.nvidia.com/nsight-systems, 2026. [24] NVIDIA Corporation. NVIDIA OpenSHMEM Library (NVSHMEM) Documentation, 2025. [25] Byungsoo Oh and Rachee Singh. Eliminating hidden serialization in multi-node megakernel communication. arXiv preprint arXiv:2605.00686, 2026. [26] Muhammad Osama, Duane Merrill, Cris Cecka, Michael Garland, and John D Owens. Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the gpu. In Proceedings of the 28th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, pages 429–431, 2023. [27] PyTorch. Pytorch symmetric memory. https://docs.pytorch.org/docs/2.14/ symmetric_memory.html. [28] Xinwei Qiang, Yue Guan, Zhengding Hu, Keren Zhou, Yufei Ding, and Adnan Aziz. Syncopate: Efficient { Multi-GPU } { AI } kernels via automatic { ChunkCentric } { Compute-Communication } overlap. In 20th USENIX Symposium on Operating Systems Design and Implementation (OSDI 26), pages 331–347, 2026. [29] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019. [30] Benjamin Spector, Jordan Juravsky, Stuart Sul, Owen Dugan, Dylan Lim, Dan Fu, Simran Arora, and Chris Ré. Look ma, no bubbles! designing a low-latency megakernel for llama-1b. https://hazyresearch.stanford.edu/blog/2025-05-27no-bubbles, May 2025. Hazy Research Blog. [31] Benjamin Spector, Jordan Juravsky, Stuart Sul, Dylan Lim, Owen Dugan, Simran Arora, and Chris Ré. We bought the whole gpu, so we’re damn well going to use the whole gpu. https://hazyresearch.stanford.edu/blog/2025-09-28-tp-llamamain, September 2025. Hazy Research Blog. [32] Stuart H Sul, Simran Arora, Benjamin F Spector, and Christopher Ré. Parallelkittens: Systematic and practical simplification of multi-gpu ai kernels. Proceedings of Machine Learning and Systems, 8:1076–1089, 2026. [33] Stuart H Sul, Nash Brown, Henry Wildermuth, William Lin, Federico Cassano, and Christopher Ré. Mixture-of-kittens: Moe megakernel for nvl72s. arXiv preprint arXiv:2609.36070, 2026. [34] Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, et al. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, pages 93–106, 2022. 10

MegaOps : A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

A

Automatic Memory Management Details

Tracking peer progress. Before reusing a symmetric memory tensor, all peers must have passed the 𝑁 -kernel mark to ensure it is safe to reuse globally. To track peer progress deterministically and without involving the host, we ask the user to mark kernels that produce a symmetric memory tensor, and kernels that consume data of such a tensor where the consumed data requires the contribution of all peers. These markers expose the dependency between local and peer kernels, allowing us to track peer progress, because the completion of such a consumer kernel implies the successful launch (and at least partial execution) of the peer producer kernels. Combined with the 𝑁 -kernel limit, we know the peers have progressed at least up to the (𝑁 − 1)-th kernel before the producer kernel. We thus add a second tensor reference to estimate the peer’s progress using the above method. Deterministic allocation. Finally, to guarantee determinism, we develop a preliminary approach where we use a dedicated memory pool for only symmetric tensors. The pool only reuses when the allocation size matches exactly the tensor in the pool, and returns the earliest free tensor for simplicity and determinism. We assume a SPMD-style programming where all peers run the same code to remove sources of inconsistencies in symmetric memory allocation. However, Python’s garbage collection (GC) may still introduce inconsistent memory release timings. We address it by asking users to manually run GC periodically (e.g., once every layer) so that all peers run at synchronized timings. To reduce frequent GC overhead, we recommend that the user freeze the currently tracked objects via gc.freeze() before running the model forward so that later GC only scans new objects created during model forward. This freeze technique is applied in our dense model forward evaluation (§5.1).

B

def TP_MLP(x, S_x): # S_x: *row-wise plain tiling* 3 if AG_RS: 4 S_pre = ... # *row-wise plain tiling* 5 x_pre = allgather(x, S_x, S_pre, ...) 6 else: # no-op, pass through 7 S_pre = S_x 8 x_pre = x 9 x_d, S_d = MLPCompute(x_pre, S_pre) 10 if AG_RS: 11 S_mlp = ... # *row-wise plain tiling* 12 x_mlp = reducescatter(x_d,S_d,S_mlp,...) 13 else: 14 S_mlp = ... # *row-wise plain tiling* 15 x_mlp = allreduce(x_d, S_d, S_mlp, ...) 16 return x_mlp, S_mlp 1 2

17

def MLPCompute(x, S_x): # shared code path # S_x: *row-wise plain tiling* 20 # gate and up projection 21 S_u = ... 22 x_u = matmul(x, S_x, S_u, ...) 23 # silu_mul 24 ... 25 # down projection 26 S_d = ... # *row-wise plain tiling* 27 x_d = matmul(x_silu, S_silu, S_d, ...) 28 return x_d, S_d Listing 4: Illustrative example showing composability enables allgather-reducescatter and allreduce to share the code path. Semaphore mappings and targets are omitted. 18 19

Implementation Details

For example, as we use two streams for scheduling subkernels, the second subkernel on the second stream may run before the synchronization signal kernel of the first subkernel, and serializing and delaying the signal until the end of the second subkernel. Using a different backend addresses the issue, and we hypothesize that it is because stream operations use complementary hardware resources to ordinary kernels.

Supported kernels. We integrate most of our computation kernels from existing open-source projects, including GeMM, normalization, elementwise, and attention kernels from TritonDistributed, and the unified attention kernel from vLLM. We added prefetch and Stream-K optimizations on top of TritonDistributed’s GeMM kernel. Communication kernels are implemented on our own, including allgather, reducescatter, distributed transpose [31] and allreduce kernels, which rely on peer memory access or multicast memory. Closed-source kernels include GeMM kernels from cuBLASLt, attention kernels from cuDNN and FlashAttention, and copy-enginebased kernels for allgather and reducescatter.2

C Extended Evaluation C.1 Extended Evaluation Setup Baselines details. We compare against the serving framework vLLM (v0.31.0), TritonDistributed (commit hash 12c2389), a distributed compiler designed for computation-communication overlapping, megakernel compilers MPK (commit hash 8b981a4) and MegaTritonKernel [38] (referred to as MegaTD) from TritonDistributed, megakernel programming framework VDCores (commit hash fd75b8d), and specialized handwritten megakernels ThunderKittens including both the single-GPU low-latency and the 8GPU high-throughput versions. For clarity, we omit MegaTD, VDCores, and the single-GPU ThunderKittens in prefill plots as they only support decode. MPK is also omitted in prefill as it is over an order of magnitude slower than ours or vLLM. In decode, we omit TritonDistributed as its AG-RS parallel scheme incurs high

Optimization for closed-source kernel scheduling. Synchronization kernels around subkernels use a complementary backend against the subkernel’s to avoid interference and achieve more timely synchronization. When the subkernel uses copy engines, we use ordinary kernels; if subkernels are ordinary kernels, we use stream memory operations [19], which is feasible as the synchronization is simply synchronizing at one variable, which is supported by stream operations. We empirically observe that using the same backend for both will contend for the same resources and get serialized. 2 FlashAttention is open-source but in CUDA, and we only support Triton currently. We thus treat it as closed-source for integration.

11

Jinkun Lin, Byungsoo Oh, and Rachee Singh

Table 2: Prefill latency (ms) with and without CTA remapping on 8×H100, using the AG-RS parallelization scheme and 1024token prompts.

overhead and makes it much slower than others. For accurate comparison, we extract the forward function vLLM uses and record its execution as a CUDA graph and wire it into the same token greedy sampler we implement to eliminate overheads caused by host and framework processing differences, which are observed to be up to 17% and 400% for decode and prefill on small models, respectively. For communication-computation overlap we compare against AsyncTP [35] and TritonDistributed.

C.6

Model

BS

Remap

No remap (slowdown)

Llama-3.1 -70B

8 16

258.7 ms 503.2 ms

303.3 ms (1.17 × ) 572.3 ms (1.14 × )

Qwen3 -8B

8 16

41.4 ms 67.2 ms

49.4 ms (1.19 × ) 88.9 ms (1.32 × )

CTA Remapping Ablation

Data tiling implementation details. Most intermediate tensors use a row-wise data tiling scheme, where each data tile covers 128 consecutive rows across the full hidden dimension, regardless of the producer kernel’s tile size. There are a few exceptions in the AR scheme, where we use two-dimensional data tiles so that consumers can start on partially produced columns, including the output of the gate and up projection, silu activation, and the self-attention output.

Table 2 shows the multi-GPU prefill latency with and without CTA remapping (§3.1) in the AG-RS scheme. Without remapping, GeMM CTAs are launched in their default order that does not match the all-gather’s order, causing delays. The result shows that the prefill latency is increased by 1.14×–1.17× on Llama-3.1-70B and 1.19×– 1.32× on Qwen3-8B without remapping, confirming the benefit of remapping.

C.2

C.7

Short-context Decode on A100 and L40

Fig. 15 shows the GeMM-reducescatter overlap benchmark from TritonDistributed, which runs the MLP down projection followed by reducescatter under tensor parallelism, using the same models for the weight shapes as the allgather-GeMM benchmark (§5.3). For M ≥ 2K tokens, the MegaOps framework performs on par with the best baselines on all models except Llama-3.1-8B, at 0.97×–1.07× of their throughput. At smaller shapes, the MegaOps framework falls behind TritonDistributed, reaching 0.62×–0.83× of its throughput on Llama-3.1-8B at M ≤ 4K and 0.85×–0.87× on Llama-3.1-70B and Qwen2-72B at M = 1K. This is because TritonDistributed switches to a different GeMM kernel that directly stores result to peer memory, while the MegaOps framework always uses the copy engine regardless of shapes, which has higher overhead in smaller shapes. We could implement a kernel-based reducescatter to replace copy engine similar to our allgather implementation. We leave it for future work.

Fig. 9 shows the short-context decode throughput on A100 and L40. The MegaOps framework achieves 1.23× (1.10×–1.48×) and 1.15× (1.00×–1.35×) higher throughput than vLLM on A100 and L40, respectively.

C.3

Short-Context Decode MBU Analysis

We use the model weights and KV cache size to estimate the total memory read traffic volume, and derive the bandwidth by dividing with execution time. We then measure the maximum achievable bandwidth using nvbandwidth on the same GPU, and calculate the memory bandwidth utilization (MBU) as the ratio of the two. Fig. 10 shows the MBU of short decode on A100 and L40. vLLM achieves a high MBU on larger models, leaving little room for MegaOps to improve. L40 has a lower achievable memory bandwidth, and thus the MegaOps framework has lower speedups on L40 than A100 in general.

C.8 C.4

Long-context Decode and Prefill

Trace Analysis of Overlap Benefits

To understand where the MegaOps framework gains and loses against non-overlapped executions and megakernels, we profile one short-context decode step at BS = 8 on H100 with intra-kernel tracing for MegaOps, MegaOps without overlap, and VDCores, and attribute time to the earliest finishing operator to account for overlap (Fig. 16). In other words, a later kernel does not count the overlapped portion with earlier kernels, but only the exposed portion after all prior kernels finish. Compared to the no-overlap version, overlap removes most of the exposed time of the small operators between GeMMs, such as SiLU and RMSNorms, and prefetching reduces the GeMM time as the loading is overlapped with prior kernels, most notably for the output projection. Compared to VDCores, the MegaOps framework exposes on par or less time on the small operators, while the GeMMs are mixed. On the other hand, our attention is faster on both models, which on Llama-3.2-1B makes MegaOps faster overall. This result shows that the MegaOps framework can achieve similar performance optimizations as a megakernel. It also shows aside from overlapping, achieving high performance for individual kernels is also important. Our system

Fig. 11 shows single-GPU decode throughput with a 1024-token prompt and output length on H100 and A100. ThunderKittens shows lower performance than vLLM while the MegaOps framework consistently outperforms vLLM on both H100 and A100, achieving on average 1.19× (1.07×–1.40×) and 1.17× (1.03×–1.38×) higher throughput than vLLM on H100 and A100, respectively.

C.5

GeMM-Reducescatter Overlap

Multi-GPU Results on 4×A100

Fig. 13 and Fig. 14 show the long-context decode and prefill throughput on 4×A100 with tensor parallelism, which follow similar trends as on 8×H100 (§5.2). For decode, the MegaOps framework outperforms vLLM by 1.15× (1.10×–1.20×) on average. For prefill, the MegaOps framework outperforms vLLM by 1.12× (1.00×–1.25×) and the non-overlapped version by 1.21×. The best parallelization scheme also depends on the batch size: on Qwen3-8B, at BS = 1, AG-RS reaches only 0.88× of AR’s throughput, while at BS = 16 it is 1.12× faster than AR. 12

MegaOps : A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

A100 L40

Normalized throughput Normalized throughput (vs. vLLM) (vs. vLLM)

Ours

Qwen3-0.6B 1.5

1.15×

1.24×

1.19×

1.17×

1.14×

1.0 0.5 0.0

BS=1 BS=2 BS=4 BS=8 BS=16

1.06×

1.14×

1.14×

1.11×

1.11×

1.09×

1.07×

1.07×

1.08×

1.09×

0.5

0.5

0.0

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

1.0

1.07×

1.07×

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

1.12×

1.10×

Qwen3-8B 1.06×

1.08×

1.05×

BS=1 BS=2 BS=4 BS=8 BS=16

BS=1 BS=2 BS=4 BS=8 BS=16

1.08×

0.0

1.10×

1.07×

1.08×

1.05×

1.04×

1.04×

BS=1 BS=2 BS=4 BS=8 BS=16

Qwen3-8B 1.05×

1.0

1.0

0.5

0.5

0.0

1.05×

1.0

Qwen3-1.7B

1.06×

0.5

0.5

1.11×

0.5

Llama-3.2-1B 1.06×

MegaTD

Qwen3-1.7B 1.0

1.10×

1.0

0.0

MPK

1.0

Qwen3-0.6B 1.5

Ours (no overlap)

Llama-3.2-1B

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

1.00×

1.03×

1.04×

1.04×

1.03×

BS=1 BS=2 BS=4 BS=8 BS=16

Figure 9: Short decode throughput on A100 (top) and L40 (bottom) using 128 prompt tokens and 64 output tokens. Numbers above bars indicate the speedup of “Ours” over the best of the other systems. VDCores and ThunderKittens do not support GPUs prior to Hopper.

H100

MBU (% of 3215 GB/s measured)

Ours 100

Qwen3-0.6B

100

A100

MBU (% of 1488 GB/s measured)

Qwen3-1.7B

100

100

80

80

80

60

60

60

60

40

40

40

40

20

20

20

20

0

BS=1 BS=2 BS=4 BS=8 BS=16

0

BS=1 BS=2 BS=4 BS=8 BS=16

100

Qwen3-0.6B

100

0

BS=1 BS=2 BS=4 BS=8 BS=16

Ours (no overlap)

Llama-3.2-1B

100

80

80

80

60

60

60

60

40

40

40

40

20

20

20

20

0

BS=1 BS=2 BS=4 BS=8 BS=16

0

BS=1 BS=2 BS=4 BS=8 BS=16

Qwen3-0.6B

0

BS=1 BS=2 BS=4 BS=8 BS=16

Ours (no overlap)

Llama-3.2-1B

0

Qwen3-1.7B 100

100

80

80

80

80

60

60

60

60

40

40

40

40

20

20

20

20

0

BS=1 BS=2 BS=4 BS=8 BS=16

0

Qwen3-8B

BS=1 BS=2 BS=4 BS=8 BS=16

Qwen3-8B

100

BS=1 BS=2 BS=4 BS=8 BS=16

BS=1 BS=2 BS=4 BS=8 BS=16

vLLM

100

0

Qwen3-8B

vLLM

Qwen3-1.7B

100

0

80

Ours

L40

vLLM

80

Ours

MBU (% of 734 GB/s measured)

Ours (no overlap)

Llama-3.2-1B

BS=1 BS=2 BS=4 BS=8 BS=16

0

BS=1 BS=2 BS=4 BS=8 BS=16

Figure 10: Short-context decode MBU analysis. The MBU is calculated against the maximum achievable bandwidth measured using nvbandwidth on the same GPU, instead of the theoretical maximum from hardware specifications.

C.9

can reuse existing kernels and easily switch between them, allowing kernels to be independently evolved and optimized. This makes achieving high performance for individual kernel much easier than in a megakernel, where the kernel is often coupled with the rest of the system and reusing other kernels is difficult.

Trace Breakdown

To understand where the multi-GPU gains come from, we profile the long-context prefill and decode on 8×H100 and attribute the time to the operator that is exposed, as in Appendix C.8 (Fig. 16). Prefill. Fig. 17 shows the prefill breakdown. We focus on analyzing AG-RS. As shown, the gain primarily comes from better overlap 13

Jinkun Lin, Byungsoo Oh, and Rachee Singh

H100

Normalized throughput (vs. vLLM)

Ours

Qwen3-0.6B 1.5 1.31×

1.38×

1.27×

1.22×

1.18× 1.15×

1.0 0.5 0.0

Ours (no overlap)

BS=1 BS=2 BS=4 BS=8 BS=16

1.18×

1.17×

A100

Normalized throughput (vs. vLLM)

1.17×

1.13×

1.14×

1.21×

1.08×

0.5

0.5

0.0

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

Ours (no overlap)

1.08×

0.5 BS=1 BS=2 BS=4 BS=8 BS=16

1.06×

1.06×

1.00×

1.06×

1.18×

Qwen3-8B 1.12×

1.09×

1.10×

1.0

0.5

0.5 BS=1 BS=2 BS=4 BS=8 BS=16

0.0

1.08×

1.07×

0.5

BS=1 BS=2 BS=4 BS=8 BS=16 MPK

1.08×

1.03×

1.10×

1.0

0.0

BS=1 BS=2 BS=4 BS=8 BS=16

MegaTD

Qwen3-1.7B

1.0

0.0

1.19×

1.19×

Llama-3.2-1B

1.0

0.0

1.16×

ThunderKittens

Qwen3-1.7B 1.0

Ours

1.19×

MegaTD

1.0

Qwen3-0.6B 1.5

MPK

Llama-3.2-1B

1.10×

1.07×

1.06×

Qwen3-8B 1.07×

1.05×

1.05×

1.04×

1.05×

1.0

1.01×

0.5

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

BS=1 BS=2 BS=4 BS=8 BS=16

Figure 11: Long decode throughput using 1024 prompt tokens and 1024 output tokens. VDcores crashes due to KV cache index integer overflow.

H100

Normalized throughput (vs. vLLM)

Ours

Qwen3-0.6B 1.0

0.94×

0.90×

0.95×

0.93×

1.01×

1.0

0.90×

0.92×

0.93×

0.96×

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

BS=1 BS=2 BS=4 BS=8 BS=16

A100

Normalized throughput (vs. vLLM)

Ours

0.93×

0.97×

1.04×

1.02×

1.01×

1.0

0.0

0.91×

0.98×

Qwen3-8B 0.99×

0.0

0.97×

1.02×

1.00×

1.00×

1.0

0.97×

0.93×

0.95×

1.00×

1.00×

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

BS=1 BS=2 BS=4 BS=8 BS=16

Ours (no overlap)

Qwen3-1.7B 0.99×

1.0

0.96×

1.01×

0.98×

1.00×

Qwen3-8B 0.99×

0.5

0.5

BS=1 BS=2 BS=4 BS=8 BS=16

0.95×

0.5

Llama-3.2-1B

0.5 0.0

1.0

0.93×

0.5

Qwen3-0.6B 1.0

Qwen3-1.7B 0.98×

0.5

0.5 0.0

Ours (no overlap)

Llama-3.2-1B

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

1.0

0.96×

1.00×

1.01×

1.03×

1.02×

0.5

BS=1 BS=2 BS=4 BS=8 BS=16

0.0

BS=1 BS=2 BS=4 BS=8 BS=16

Figure 12: Prefill throughput using 1024 prompt tokens. MPK

Llama-3.1-70B

1.10× 1.09× 1.10× 1.07× 1.05×

1.04× 1.04×

1.0

1.0

0.5

0.5

0.0

BS=1 BS=2 BS=4 BS=8BS=16

0.0

Ours

4xA100

Ours (no overlap)

Qwen3-8B

BS=1 BS=2

Figure 13: Long decode throughput on 4×A100. “BS” denotes global batch size. MPK does not support Llama-3.1-70B, and BS>2 runs out-of-memory for all systems on Llama-3.1-70B.

Normalized throughput (vs. vLLM)

4xA100

Normalized throughput (vs. vLLM)

Ours

Ours (no overlap)

TritonDistributed

Qwen3-8B 1.11×

1.16× 1.13× 1.18× 1.17×

Llama-3.1-70B 1.0

1.01× 1.00×

1.0 0.5

0.5 0.0

BS=1 BS=2 BS=4 BS=8BS=16

0.0

BS=1 BS=2

Figure 14: Prefill throughput on 4×A100. BS>2 runs out-ofmemory for all systems on Llama-3.1-70B. Decode. Fig. 18 shows the decode breakdown. Unlike prefill, the gain in decode mainly comes from reduced GeMM and small operator times, consistent with the observation in Fig. 16. In this small shape GeMMs finish in one wave, leaving no room to hide the AR communication.

between computation and communication. Compared to AG-RS without overlap and vLLM, which do not overlap communication with GeMMs, AG-RS with overlap hides most of its communication (observe the large reduction in cyan bars from “B” to “A”), with only a small amount of impact on the GeMM performance. At large BS, this overlap pays off and hides more communication than the AR scheme, making AG-RS the fastest system. 14

MegaOps : A Kernel Overlapping Framework with PyTorch-like Modularity and Megakernel Performance

Ours

8xH100

Normalized throughput (vs. Serial PyTorch)

LLaMA-3.1-8B 1.02×

1.5 1.0

0.62×

0.69×

0.83×

1K

2K

1.5

1.02×

0.99×

4K 8K M (tokens)

1.01×

16K

1.5 0.98×

1.01×

1.0 0.87×

1.0

0.5

0.5

0.0

1K

2K

4K 8K M (tokens)

AsyncTP

Qwen2-72B

1.06×

0.5 0.0

TritonDistributed

LLaMA-3.1-70B

16K

0.0

0.85×

1K

0.99×

GPT-3-175B 0.97×

1.07×

2K

4K 8K M (tokens)

16K

1.01×

1.02×

1.00×

LLaMA-3.1-405B 0.99×

0.98×

1.0 0.96×

1.0

0.5

0.5

0.0

1K

2K

4K 8K M (tokens)

16K

0.0

1.01×

1.00×

0.97×

0.99×

1K

2K

4K 8K M (tokens)

16K

Figure 15: GeMM-reducescatter overlap on 8×H100. Numbers above bars indicate the speedup of “Ours” over the best of the other systems.

Qwen3-1.7B Llama-3.2-1B BS=8 BS=8 1.4

Exposed time (ms)

2.0

1.2 1.0

1.5

0.8 1.0

0.6 0.4

0.5 0.0

0.2

O N V

0.0

Sampling LM head Final RMSNorm Down proj Output zeroing SiLU-mul Gate/up proj RMSNorm (pre-MLP) O proj Attention QK-norm+RoPE+KV write QKV proj RMSNorm (pre-attn) Embedding

O N V

Figure 16: Exposed time per operator of one short-context decode step at BS = 8 on H100. O, N, V denote MegaOps, MegaOps without overlap, and VDCores, respectively.

15

Jinkun Lin, Byungsoo Oh, and Rachee Singh

Qwen3-8B

Llama-3.1-70B 700

Exposed prefill time (ms)

100

600 80

500

60

400 300

40

200 20 0

100 ON A B V BS=1

ON A B V BS=2

ON A B V BS=4

ON A B V BS=8

0

ON A B V BS=16

ON A B V BS=1

ON A B V BS=2

ON A B V BS=4

ON A B V BS=8

ON A B V BS=16

O = Ours, N = Ours (no overlap), A = Ours (AG-RS), B = Ours (no overlap, AG-RS), V = vLLM Embedding RMSNorm / RoPE / SiLU-mul / Stream-K zeroing LM head QKV / O / Gate-up / Down proj Comm (AR / AG / RS / transpose) Sampling Attention + combine Figure 17: Exposed time per operator group of the 1024-token prefill step on 8×H100. O, N, A, B, V denote MegaOps, MegaOps without overlap, MegaOps with AG-RS, MegaOps with AG-RS and without overlap, and vLLM, respectively.

Exposed time per decode step (ms)

Qwen3-8B

Llama-3.1-70B 16

4

14 12

3

10 8

2

6 4

1

2 0

O N A B V BS=1

O N A B V BS=2

O N A B V BS=4

O N A B V BS=8

0

O N A B V BS=16

O N A B V BS=1

O N A B V BS=2

O N A B V BS=4

O N A B V BS=8

O N A B V BS=16

O = Ours, N = Ours (no overlap), A = Ours (AG-RS), B = Ours (no overlap, AG-RS), V = vLLM Embedding RMSNorm / RoPE / SiLU-mul / Stream-K zeroing LM head QKV / O / Gate-up / Down proj Comm (AR / AG / RS / transpose) Sampling Attention + combine Figure 18: Exposed time per operator group of one long-context decode step on 8×H100. O, N, A, B, V denote MegaOps, MegaOps without overlap, MegaOps with AG-RS, MegaOps with AG-RS and without overlap, and vLLM, respectively. vLLM fuses the allreduce with the following RMSNorm at BS ≤ 8, which is counted as communication.

16

Related documents

• How Much of a Real Workload Can LLM-Generated GPU Kernels Actually Reach?#1035028
• How Much of a Real Workload Can LLM-Generated GPU Kernels Actually Reach?#1006828
• KEET: Explaining Performance of GPU Kernels Using LLM Agents#168286
• Exo-GPU: Safe, Imperative, User-schedulable Programming for Tensor Cores#1110750
• PipeLink: Contract-Guided Linking of GPU Kernel Components#1473574