Nvfp4 A16 GEMM Plugin

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

TensorRT V3 plugin for a dense FP16-activation / Marlin-packed NVFP4 (W4A16) GEMM.

This is a thin wrapper over the MoE Marlin FP16xE2M1 kernel (moeNvfp4A16MarlinGemm): it drives the kernel with a single expert and top_k = 1 so that every input row maps to expert 0. The weights, block scales, and global scale carry a leading expert dimension of 1 to reuse the MoE kernel’s shape contract unchanged.

gemm_n is the Marlin-padded output dimension (a multiple of 128); the caller slices the logical width downstream. max_m is the optimization-profile token capacity (batch * sequence) used to size the routing/Marlin workspace.

Public Functions

Nvfp4A16GemmPlugin(

std::string const &name,

int32_t gemmN,

int32_t gemmK,

int32_t maxM

)

Nvfp4A16GemmPlugin(

std::string const &name,

nvinfer1::PluginFieldCollection const *fc

)

Nvfp4A16GemmPlugin() = delete

Nvfp4A16GemmPlugin(Nvfp4A16GemmPlugin const&) = delete

Nvfp4A16GemmPlugin &operator=(Nvfp4A16GemmPlugin const&) = delete

~Nvfp4A16GemmPlugin() noexcept override

nvinfer1::IPluginCapability *getCapabilityInterface(

nvinfer1::PluginCapabilityType type

) noexcept override

nvinfer1::IPluginV3 *clone() noexcept override

char const *getPluginName() const noexcept override

char const *getPluginVersion() const noexcept override

char const *getPluginNamespace() const noexcept override

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

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

void setPluginNamespace(char const *pluginNamespace) noexcept

class Nvfp4A16GemmPluginCreator : public nvinfer1::IPluginCreatorV3One

Creator for Nvfp4A16GemmPlugin.

Public Functions

Nvfp4A16GemmPluginCreator()

~Nvfp4A16GemmPluginCreator() override = default

char const *getPluginName() const noexcept override

char const *getPluginVersion() const noexcept override

nvinfer1::PluginFieldCollection const *getFieldNames(

) noexcept override

char const *getPluginNamespace() const noexcept override

void setPluginNamespace(char const *pluginNamespace) noexcept

nvinfer1::IPluginV3 *createPlugin(

char const *name,

nvinfer1::PluginFieldCollection const *fc,

nvinfer1::TensorRTPhase phase

) noexcept override