All Reduce Plugin

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

TensorRT plugin for tensor-parallel all-reduce.

This plugin sums 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 tensor-parallel ranks.

The plugin is a passthrough when tpSize=1 (no tensor parallelism). For TP, it uses an optional accelerated execution path when supported and otherwise falls back to the required NCCL path registered by the runtime.

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

struct NcclAllReducePathRegistration

One device’s NCCL resources used by plugin execution paths.

Public Members

void *communicator = {nullptr}

void *allReduceFunction = {nullptr}

struct AllReducePathRegistrations

Consistent snapshot of every registered execution path for one device.

Public Members

NcclAllReducePathRegistration nccl = {}

bool trt_edgellm::plugins::registerNcclAllReducePath(

int32_t deviceId,

void *ncclComm,

void *ncclAllReduceFunction

) noexcept

Register NCCL resources for one CUDA device.

bool trt_edgellm::plugins::unregisterNcclAllReducePath(

int32_t deviceId,

void *expectedNcclComm

) noexcept

Remove NCCL resources only if the expected communicator is still registered.

AllReducePathRegistrations trt_edgellm::plugins::snapshotAllReducePathRegistrationsForDevice(

int32_t deviceId

) noexcept

Snapshot all available path registrations while holding the registry lock once.