All Reduce Plugin#

class AllReducePlugin : public nvinfer1::IPluginV3, public nvinfer1::IPluginV3OneCore, public nvinfer1::IPluginV3OneBuild, public nvinfer1::IPluginV3OneRuntime#

TensorRT plugin for NCCL all-reduce in tensor parallel inference.

This plugin performs NCCL all-reduce (sum) on intermediate activations during TensorRT engine execution. It is inserted after row-parallel linear layers (attention output projection and MLP down projection) to combine partial results from different tensor parallel ranks.

The plugin is a passthrough when tpSize=1 (no tensor parallelism). It uses communicator handles registered by the runtime plugin communication registry.

IPluginV3 Methods

nvinfer1::IPluginCapability *getCapabilityInterface(
nvinfer1::PluginCapabilityType type,
) noexcept override#
nvinfer1::IPluginV3 *clone() noexcept override#

IPluginV3OneCore Methods

char const *getPluginName() const noexcept override#
char const *getPluginNamespace() const noexcept override#
char const *getPluginVersion() const noexcept override#

IPluginV3OneBuild Methods

int32_t getNbOutputs() const noexcept override#
int32_t getOutputDataTypes(
nvinfer1::DataType *outputTypes,
int32_t nbOutputs,
nvinfer1::DataType const *inputTypes,
int32_t nbInputs,
) const noexcept override#
int32_t getOutputShapes(
nvinfer1::DimsExprs const *inputs,
int32_t nbInputs,
nvinfer1::DimsExprs const *shapeInputs,
int32_t nbShapeInputs,
nvinfer1::DimsExprs *outputs,
int32_t nbOutputs,
nvinfer1::IExprBuilder &exprBuilder,
) noexcept override#
bool supportsFormatCombination(
int32_t pos,
nvinfer1::DynamicPluginTensorDesc const *inOut,
int32_t nbInputs,
int32_t nbOutputs,
) noexcept override#
int32_t configurePlugin(
nvinfer1::DynamicPluginTensorDesc const *in,
int32_t nbInputs,
nvinfer1::DynamicPluginTensorDesc const *out,
int32_t nbOutputs,
) noexcept override#
size_t getWorkspaceSize(
nvinfer1::DynamicPluginTensorDesc const *inputs,
int32_t nbInputs,
nvinfer1::DynamicPluginTensorDesc const *outputs,
int32_t nbOutputs,
) const noexcept override#

IPluginV3OneRuntime Methods

int32_t enqueue(
nvinfer1::PluginTensorDesc const *inputDesc,
nvinfer1::PluginTensorDesc const *outputDesc,
void const *const *inputs,
void *const *outputs,
void *workspace,
cudaStream_t stream,
) noexcept override#
int32_t onShapeChange(
nvinfer1::PluginTensorDesc const *in,
int32_t nbInputs,
nvinfer1::PluginTensorDesc const *out,
int32_t nbOutputs,
) noexcept override#
nvinfer1::IPluginV3 *attachToContext(
nvinfer1::IPluginResourceContext *context,
) noexcept override#
nvinfer1::PluginFieldCollection const *getFieldsToSerialize(
) noexcept override#

Public Functions

AllReducePlugin(std::string const &name, int32_t tpSize)#

Constructor with TP configuration.

Parameters:

  • name – [in] Plugin instance name

  • tpSize – [in] Tensor parallel world size

AllReducePlugin(
std::string const &name,
nvinfer1::PluginFieldCollection const *fc,
)#

Constructor from TensorRT plugin fields.

Parameters:

  • name – [in] Plugin instance name

  • fc – [in] Plugin field collection

AllReducePlugin() = delete#
AllReducePlugin(AllReducePlugin const&) = delete#
~AllReducePlugin() override#
void setPluginNamespace(char const *pluginNamespace) noexcept#
class AllReducePluginCreator : public nvinfer1::IPluginCreatorV3One#

Factory class for creating AllReducePlugin instances.

Public Functions

AllReducePluginCreator()#
~AllReducePluginCreator() override = default#
char const *getPluginName() const noexcept override#
nvinfer1::PluginFieldCollection const *getFieldNames(
) noexcept override#
void setPluginNamespace(char const *pluginNamespace) noexcept#
char const *getPluginNamespace() const noexcept override#
char const *getPluginVersion() const noexcept override#
nvinfer1::IPluginV3 *createPlugin(
char const *name,
nvinfer1::PluginFieldCollection const *fc,
nvinfer1::TensorRTPhase phase,
) noexcept override#
void trt_edgellm::plugins::registerNcclCommForAllReducePlugin(
int deviceId,
void *ncclComm,
void *ncclAllReduceFunc,
) noexcept#

Register NCCL communicator for use by AllReducePlugin instances.

Must be called by the runtime after NCCL initialization and before any TRT engine execution that contains AllReducePlugin nodes. This avoids linking the plugin shared library against the runtime library.

Parameters:

  • deviceId – CUDA device ID this communicator is for

  • ncclComm – NCCL communicator handle (ncclComm_t)

  • ncclAllReduceFunc – Pointer to ncclAllReduce function

void trt_edgellm::plugins::getNcclRegistrationForDevice(
int deviceId,
void **ncclComm,
void **ncclAllReduceFunc,
) noexcept#

Snapshot the NCCL communicator and AllReduce function for one CUDA device.

void *trt_edgellm::plugins::getNcclCommForDevice(
int deviceId,
) noexcept#

Get the NCCL communicator for a given CUDA device.

Returns:

ncclComm_t handle, or nullptr if not registered

void *trt_edgellm::plugins::getNcclAllReduceFunc() noexcept#

Get the registered NCCL AllReduce function pointer.