Implementation:NVIDIA TransformerEngine Comm GEMM

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

Overview

Implements communication-fused GEMM operations using cuBLASMp and NVSHMEM, enabling distributed matrix multiplications with overlapped AllGather and ReduceScatter collective communication across multiple GPUs.

Description

comm_gemm.cpp manages a NVTECommGemmCtx context that holds cuBLASMp handles, CUDA streams/events, and grid descriptors. It provides AgGemmInitMatrices and GemmRsInitMatrices to set up block-cyclic matrix descriptors for AllGather-GEMM and GEMM-ReduceScatter patterns respectively, using non-cyclic layouts to maximize communication overlap.

Key components:

• NVTECommGemmCtx: Struct holding per-context state including NCCL communicator, cuBLASMp handle, grid descriptors (column-major and row-major), matrix descriptors, and workspace allocation.
• RAII wrappers: Template-based CreateWithCudaCheck and CreateWithCublasMpCheck wrappers that provide automatic resource cleanup for CUDA and cuBLASMp handles.
• AllGather+GEMM: Gathers distributed data from all ranks, then computes the matrix multiplication with overlapped communication.
• GEMM+ReduceScatter: Computes the matrix multiplication, then distributes results across ranks with reduction.

Usage

Use when performing tensor-parallel GEMM operations in multi-GPU Transformer training where communication and computation need to be overlapped to minimize latency.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/comm_gemm/comm_gemm.cpp

Lines

1--519

Signature

struct NVTECommGemmCtx {
  int64_t nranks;
  int64_t rank;
  ncclComm_t comm;
  CudaStream stream;
  CudaEvent event;
  CublasMp cublas_mp;
  CublasMpGrid grid_col_major;
  CublasMpGrid grid_row_major;
  CublasMpMatrixDesc a_desc, b_desc, d_desc;
  CublasMpMatmulDesc matmul_desc;
  void* workspace;
  size_t workspace_size;
};

NVTECommGemmCtx* nvte_comm_gemm_ctx_create(ncclComm_t comm, int nranks, int rank);
void nvte_comm_gemm_ctx_destroy(NVTECommGemmCtx* ctx);

Import

#include "transformer_engine/comm_gemm.h"

I/O Contract

Inputs

Name	Type	Required	Description
comm	ncclComm_t	Yes	NCCL communicator for the tensor-parallel group
nranks	int	Yes	Number of ranks in the TP group
rank	int	Yes	Local rank within the TP group
a	NVTETensor	Yes	Local part of A matrix
b	NVTETensor	Yes	Local part of B matrix

Outputs

Name	Type	Description
d	NVTETensor	Local part of the result matrix D

Usage Examples

#include "transformer_engine/comm_gemm.h"

// Create context with NCCL communicator
NVTECommGemmCtx* ctx = nvte_comm_gemm_ctx_create(nccl_comm, tp_size, tp_rank);

// Perform AllGather + GEMM
nvte_all_gather_gemm(ctx, m, n, k, a, b, d, bias, pre_act_out,
                     transa, transb, grad, accumulate, 0, stream,
                     kNVTECommGemmAlgoDefault);

// Cleanup
nvte_comm_gemm_ctx_destroy(ctx);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements