Implementation:NVIDIA TransformerEngine Comm GEMM C API

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

Overview

Declares the C API for distributed (multi-GPU) matrix multiplication that fuses NCCL communication with GEMM computation, providing a TE-native binding to the cuBLASMp library for tensor-parallel patterns.

Description

comm_gemm.h defines three core distributed GEMM operations:

• nvte_all_gather_gemm: AllGather + GEMM -- gathers distributed data from all ranks then computes matrix multiplication with overlapped communication.
• nvte_gemm_reduce_scatter: GEMM + ReduceScatter -- computes matrix multiplication then distributes results across ranks with reduction.
• nvte_gemm_all_reduce: GEMM + AllReduce -- computes matrix multiplication then all-reduces the result.

An NVTECommGemmAlgoType enum controls the algorithm:

• kNVTECommGemmAlgoDefault: Default heuristic
• kNVTECommGemmAlgoSplitP2P: Split with P2P transfers
• kNVTECommGemmAlgoSplitMulticast: Split with multicast
• kNVTECommGemmAlgoAtomicP2P: Atomic with P2P
• kNVTECommGemmAlgoAtomicMulticast: Atomic with multicast

Usage

Use for tensor-parallel training where communication and computation must be overlapped for column-parallel and row-parallel linear layers.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/include/transformer_engine/comm_gemm.h

Lines

1--156

Signature

typedef struct NVTECommGemmCtx NVTECommGemmCtx;

enum NVTECommGemmAlgoType {
  kNVTECommGemmAlgoDefault = 0,
  kNVTECommGemmAlgoSplitP2P = 1,
  kNVTECommGemmAlgoSplitMulticast = 2,
  kNVTECommGemmAlgoAtomicP2P = 3,
  kNVTECommGemmAlgoAtomicMulticast = 4
};

NVTECommGemmCtx* nvte_comm_gemm_ctx_create(ncclComm_t comm, int nranks, int rank);
void nvte_comm_gemm_ctx_destroy(NVTECommGemmCtx* ctx);
void nvte_all_gather_gemm(NVTECommGemmCtx* ctx, int64_t m, int64_t n, int64_t k, ...);
void nvte_gemm_reduce_scatter(NVTECommGemmCtx* ctx, int64_t m, int64_t n, int64_t k, ...);
void nvte_gemm_all_reduce(NVTECommGemmCtx* ctx, int64_t m, int64_t n, int64_t k, ...);

Import

#include "transformer_engine/comm_gemm.h"

I/O Contract

Inputs

Name	Type	Required	Description
ctx	NVTECommGemmCtx*	Yes	Communication-GEMM context
m, n, k	int64_t	Yes	Global matrix dimensions
a, b	NVTETensor	Yes	Local matrix operands
algo	NVTECommGemmAlgoType	Yes	Algorithm selection

Outputs

Name	Type	Description
d	NVTETensor	Local part of the result matrix

Usage Examples

#include "transformer_engine/comm_gemm.h"

NVTECommGemmCtx* ctx = nvte_comm_gemm_ctx_create(nccl_comm, tp_size, tp_rank);

// Column-parallel forward: AllGather + GEMM
nvte_all_gather_gemm(ctx, m, n, k, a, b, d, bias, pre_act_out,
                     transa, transb, grad, accumulate,
                     comm_sm_count, stream, kNVTECommGemmAlgoDefault);

// Row-parallel forward: GEMM + ReduceScatter
nvte_gemm_reduce_scatter(ctx, m, n, k, a, b, d, bias, pre_act_out,
                         transa, transb, grad, accumulate,
                         comm_sm_count, stream, kNVTECommGemmAlgoDefault);

nvte_comm_gemm_ctx_destroy(ctx);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements