Implementation:NVIDIA TransformerEngine Comm GEMM Overlap API

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, Distributed_Computing
Last Updated	2026-02-07 14:00 GMT

Overview

Defines the C++ class hierarchy for overlapping NCCL collective communication with GEMM computation using the Userbuffers infrastructure, enabling pipelined tensor-parallel execution across multiple GPU SMs.

Description

comm_gemm_overlap.h provides three classes forming a hierarchy:

• CommOverlapCore: Base class with Userbuffers communicator, SM partitioning, and virtual methods for overlap strategies. Manages compute/communication CUDA streams, events, and atomic GEMM counters.
• CommOverlapBase: Adds stream-based bulk and split-pipelined reduce-scatter overlap.
• CommOverlapP2PBase: Adds P2P-based split-pipelined and atomic GEMM overlap for both AllGather and ReduceScatter.

Supported overlap algorithms via CommOverlapAlgo enum:

• BULK_OVERLAP_AG/RS: Bulk communication overlapped with GEMM
• SPLIT_PIPELINED_AG_P2P/RS/RS_P2P: Split-pipelined communication with point-to-point transfers
• ATOMIC_GEMM_RS/AG_P2P/RS_P2P: Atomic GEMM with fine-grained overlap
• EXTERNAL_BULK_OVERLAP_AG: External bulk overlap for AllGather

Usage

Used by framework bindings to create comm+GEMM overlap objects for tensor-parallel linear layers. Configure via parameters controlling CGA size, SM allocation, stream priorities, and CE usage.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h

Lines

1--327

Signature

namespace transformer_engine {

enum class CommOverlapType { RS = 0, AG = 1 };

enum class CommOverlapAlgo {
  BULK_OVERLAP_AG = 0, BULK_OVERLAP_RS = 1,
  SPLIT_PIPELINED_AG_P2P = 2, SPLIT_PIPELINED_RS = 3,
  SPLIT_PIPELINED_RS_P2P = 4, ATOMIC_GEMM_RS = 5,
  ATOMIC_GEMM_AG_P2P = 6, ATOMIC_GEMM_RS_P2P = 7,
  EXTERNAL_BULK_OVERLAP_AG = 8,
};

class CommOverlapCore { ... };
class CommOverlapBase : public CommOverlapCore { ... };
class CommOverlapP2PBase : public CommOverlapCore { ... };

}  // namespace transformer_engine

Import

#include <transformer_engine/comm_gemm_overlap.h>

I/O Contract

Inputs

Name	Type	Required	Description
A	TensorWrapper	Yes	First matrix operand for GEMM
B	TensorWrapper	Yes	Second matrix operand for GEMM
comm_type	CommOverlapType	Yes	RS (ReduceScatter) or AG (AllGather)
stream_main	cudaStream_t	Yes	Main CUDA stream

Outputs

Name	Type	Description
D	TensorWrapper	GEMM output matrix
rs_output	TensorWrapper	ReduceScatter output (for RS overlap)

Usage Examples

#include <transformer_engine/comm_gemm_overlap.h>

using namespace transformer_engine;

// Create a CommOverlapBase for RS overlap
// (typically done via Python bindings)
CommOverlapBase overlap(myrank, numranks, mylocal, numlocal,
                        mynode, numnodes, tp_size,
                        allgather_fn, barrier_fn,
                        num_splits, num_max_streams,
                        comm_cga_size, gemm_priority, comm_priority,
                        num_comm_sm, set_sm_margin, use_ce, atomic_gemm);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements