Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts

• 47% of total execution time is consumed by inter-device communication in the forward pass of several popular MoE models.Source: §1 · p. 2
• Coarse-grained overlap partitions input data into chunks, but smaller expert chunks can reduce GPU utilization and leave non-overlapping communication phases.Partitioned expert computation can take longer than the original computation, while initial and final communication phases create GPU idle time.Source: §1 · p. 2
• COMET resolves communication-computation dependencies and adaptively assigns GPU thread blocks to balance workloads within kernels.Its dependency method structures the pipeline, while adaptive assignment improves latency concealment under dynamic workloads.Source: §1 · pp. 2–3
• 1.96× speedup is achieved for typical MoE layers, while end-to-end MoE model execution gains 1.71× on average versus SOTA MoE systems.Source: §1 · p. 3
• COMET has been deployed in production clusters with over ten thousand GPUs, saving millions of GPU hours.Source: §1 · p. 3

2.1 MoE Structure

• MoE activates only selected experts for each token, using routing and feed-forward computations before gathering and reducing the topk outputs.Source: §2.1 · p. 3
• Expert parallelism distributes intact expert weights across GPUs, while tensor parallelism partitions all expert weights along the hidden dimension; deployments often combine both.Source: §2.1 · pp. 3–4

2.2 Background and Motivation

• Fine-grained overlap is motivated by the increasing communication share of larger, sparser MoE models and the limited optimization potential of coarse-grained schemes.Source: §2.2 · p. 4

2.2.1 Granularity mismatch between computation and communication

• A computation tile cannot start until all of its required tokens arrive, because runtime routing distributes those tokens irregularly across devices.Fine-grained communication and data reorganization are proposed to expose only the data each tile needs and hide preparation with computation.Source: §2.2.1 · p. 4
• Remote token I/O can block computation thread blocks, so integrating token-wise communication with tile-wise computation risks reducing kernel efficiency.Source: §2.2.1 · p. 4

2.2.2 Diverse loads of computation and communication

• COMET resolves granularity mismatch through shared-tensor decomposition and reorganization, then adapts in-kernel resource allocation for communication and computation.Source: §2.2.1 · p. 4, §2.2.2 · p. 5, Figure 3 · p. 5
• Dynamic routing produces varying expert input shapes, compute demands, and communication demands across GPUs and hardware environments.Source: §2.2.2 · p. 4

3.1 Shared Tensor Based Dependency Resolving

• Comet analyzes shared-tensor dependencies to decompose data along consumer-independent dimensions and reschedule computation for overlap.Shared tensors bridge producer and consumer operators; layer0 decomposes along M, while layer1 decomposes along N because top-k reduction creates M-dimension interdependencies.Source: §3 · p. 5, §3.1.1 · p. 6

3.1.2 How to reschedule the decomposed shared tensor?

• Rescheduled sub-tensors align with computation tiles while prioritizing producer outputs that the consumer can use immediately.Fine-grained row- or column-level decomposition alone can reduce GEMM efficiency, so Comet reorganizes sub-tensors into tiles.Source: §3.1.2 · p. 6
• Layer0 sorts tokens by source rank and prioritizes tiles containing local tokens while remote-token transfers proceed concurrently.GroupGEMM computes experts on the current rank, reducing dependence on remote data during the communication-computation pipeline.Source: §3.1.2 · p. 6
• Layer1 executes GroupGEMM operations column-wise so reduction and communication can begin after the first TN columns are produced.This rescheduling lets the consumer start before all expert computations finish, unlike sequential expert execution.Source: §3.1.2 · p. 6

3.2 Adaptive Workload Assignment

• Comet combines fused communication and computation with thread-block specialization to isolate communication effects and balance workload latencies.Adjusting thread-block allocation reduces overlap bubbles, while specialized kernels preserve computation performance by separating communication and computation resources.Source: §3.2 · p. 7, §3.2.1 · p. 7

3.2.1 Thread block specialization

• Thread-block specialization remains constrained by hardware resources because communication warps can limit bandwidth utilization and interfere with computation warps.Integrating communication and computation warps could reduce redundant global-memory accesses, but warp-count restrictions create competing resource demands.Source: §3.2.1 · p. 7

3.2 Adaptive Workload Assignment

• The optimal communication–computation resource division changes with input length and parallel configuration.For TP = 8, the optimal nc changes from 18 to 26 as M increases from 4096 to 16384.Source: §3.2.2 · p. 8

4 Implementation

• Comet optimizes MoE GEMM kernels with CUTLASS templates and register caching of row indices to reduce global-memory access costs.The row indices are otherwise accessed from global memory at each K iteration in MoE layer0.Source: §4 · p. 8
• Comet uses NVSHMEM to provide fine-grained GPU-initiated communication through a global address space spanning multiple GPUs.Compared with NCCL’s higher-level operations, NVSHMEM offers a more composable low-level API for fine-grained data access within kernels.Source: §4 · p. 8

Overall Performance

• 34.1%, 42.6%, 44.4% and 31.8% lower end-to-end latency than Megatron-Cutlass, Megatron-TE, FasterMoE and Tutel, respectively.Comet outperforms the baselines in all tested configurations.Source: §5.2 · p. 9

Detailed Evaluation on a Single MoE Layer

• 1.28× to 2.37× speedup over baselines across varying input token lengths.The advantage is especially prominent for small M, where host-side scheduling dominates and Comet moves scheduling into fused kernels.Source: §5.3 · pp. 9–10
• 86.5% of communication latency hidden by Comet, compared with 29.2% for FasterMoE and 68.6% for Tutel.Comet maintains expert computational efficiency while overlapping communication.Source: §5.3 · p. 10
• Comet maintains low latency across diverse parallelism strategies, while tensor parallelism increases latency for other supported baselines.Comet reschedules shared tensors to preserve computational efficiency and eliminates weight-switching overhead.Source: §5.3 · p. 10

Scaling to Distinct Clusters

• 1.19× to 1.46× average speedup over baselines on an L20 cluster with PCIe-connected GPUs and roughly 25 GB/s GPU-to-GPU bandwidth.The result covers a bandwidth-limited environment distinct from the H800 cluster.Source: §5.4 · p. 11

Communication optimization

• Communication optimization methods accelerate data transmission, use 2D-hierarchical all-to-all algorithms, or reduce communication volume through compression.These approaches target communication cost directly rather than primarily changing computation scheduling.Source: §6 · p. 11

Computation-communication overlapping

• FasterMoE and Tutel pipeline all-to-all communication with expert computation using fixed, manually set, or heuristically searched pipeline degrees.FasterMoE uses a pipeline degree of 2, while Tutel supports manual settings or heuristic search under a limited search space.Source: §6 · p. 12

Conclusion

• Comet uses shared-tensor dependency resolution to enable fine-grained overlap while eliminating fine-grained communication I/O bottlenecks.The design is paired with workload assignment for adaptive operator overlap and latency concealing.Source: §7 · p. 12
• 1.96× speedup for a single MoE layer and 1.71× speedup for end-to-end MoE execution compared with existing literature.These are the paper’s reported aggregate speedups.Source: §7 · p. 12