Int4 Groupwise GEMM Plugin V2#

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

TensorRT plugin for INT4 group-wise quantized (W4A16) GEMM, cuteDSL backend.

Consumes weights already in the cuteDSL fragment layout (repacked at export time, bN=128 pinned), and dispatches to the AOT-compiled CuTe DSL W4A16 GEMM variants selected by the TensorRT plugin-V3 autotuner. No runtime weight repack and no cached weight buffer.

Inputs: [0] activations [b,seq,K] fp16; [1] fragment weights [rows,512] int8 (rows = ceil(N/128)*ceil(K/64)*8); [2] scales [ceil(K/G),N] fp16. Output: [b,seq,N] fp16.

Public Functions

Int4GroupwiseGemmPluginV2(
std::string const &name,
int32_t N,
int32_t K,
int32_t groupSize
)#

Construct INT4 group-wise GEMM plugin.

Parameters:

  • name – Layer name

  • N – Output dimension (columns in weight matrix)

  • K – Input dimension (rows in weight matrix)

  • groupSize – Quantization group size

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

Construct from field collection.

Parameters:

  • name – Layer name

  • fc – Plugin field collection

Int4GroupwiseGemmPluginV2() = delete#

Deleted default constructor.

Int4GroupwiseGemmPluginV2(Int4GroupwiseGemmPluginV2 const&) = delete#

Deleted copy constructor.

~Int4GroupwiseGemmPluginV2() override#

Destructor.

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

Return the plugin capability interface for given type.

nvinfer1::IPluginV3 *clone() noexcept override#

Clone the plugin for use in another network.

Returns:

Cloned plugin instance

char const *getPluginName() const noexcept override#

Get plugin name.

Returns:

Plugin name string

char const *getPluginVersion() const noexcept override#

Get plugin version.

Returns:

Version string

char const *getPluginNamespace() const noexcept override#

Get plugin namespace.

Returns:

Namespace string

int32_t getNbOutputs() const noexcept override#

Get number of output tensors.

Returns:

Number of outputs (1)

int32_t getOutputDataTypes(
nvinfer1::DataType *outputTypes,
int32_t nbOutputs,
nvinfer1::DataType const*,
int32_t
) const noexcept override#

Get output tensor data types.

Parameters:

  • outputTypes – Output array for data types

  • nbOutputs – Number of outputs

  • inputTypes – Input data types (not used)

  • nbInputs – Number of inputs (not used)

Returns:

0 on success, non-zero on error

int32_t getOutputShapes(
nvinfer1::DimsExprs const *inputs,
int32_t nbInputs,
nvinfer1::DimsExprs const*,
int32_t,
nvinfer1::DimsExprs *outputs,
int32_t nbOutputs,
nvinfer1::IExprBuilder &exprBuilder
) noexcept override#

Get output tensor shapes.

Parameters:

  • inputs – Input dimensions

  • nbInputs – Number of inputs

  • shapeInputs – Shape tensor inputs (not used)

  • nbShapeInputs – Number of shape inputs (not used)

  • outputs – Output dimensions

  • nbOutputs – Number of outputs

  • exprBuilder – Expression builder for dynamic shapes

Returns:

0 on success, non-zero on error

bool supportsFormatCombination(
int32_t pos,
nvinfer1::DynamicPluginTensorDesc const *inOut,
int32_t nbInputs,
int32_t nbOutputs
) noexcept override#

Check if format combination is supported.

Parameters:

  • pos – Position in input/output array

  • inOut – Input and output tensor descriptors

  • nbInputs – Number of inputs

  • nbOutputs – Number of outputs

Returns:

True if supported

int32_t configurePlugin(
nvinfer1::DynamicPluginTensorDesc const *in,
int32_t nbInputs,
nvinfer1::DynamicPluginTensorDesc const*,
int32_t
) noexcept override#

Configure plugin with tensor descriptions.

Captures the profile’s optimum token count M for the autotune tile-pruning heuristic (see getValidTactics).

Parameters:

  • in – Input tensor descriptors (activation opt dims read for autotune)

  • nbInputs – Number of inputs

  • out – Output tensor descriptors (not used)

  • nbOutputs – Number of outputs (not used)

Returns:

0 on success, non-zero on error

size_t getWorkspaceSize(
nvinfer1::DynamicPluginTensorDesc const *inputs,
int32_t,
nvinfer1::DynamicPluginTensorDesc const*,
int32_t
) const noexcept override#

Get workspace size required for execution.

Parameters:

  • inputs – Input tensor descriptors (max dims used to size the lock buffer)

  • nbInputs – Number of inputs (not used)

  • outputs – Output tensor descriptors (not used)

  • nbOutputs – Number of outputs (not used)

Returns:

Bytes for the serial split-K int32 lock buffer (0 when no split-K tactic applies)

int32_t enqueue(
nvinfer1::PluginTensorDesc const *inputDesc,
nvinfer1::PluginTensorDesc const*,
void const *const *inputs,
void *const *outputs,
void *workspace,
cudaStream_t stream
) noexcept override#

Execute the plugin.

Parameters:

  • inputDesc – Input tensor descriptors

  • outputDesc – Output tensor descriptors

  • inputs – Input tensor pointers

  • outputs – Output tensor pointers

  • workspace – Workspace pointer

  • stream – CUDA stream

Returns:

0 on success, non-zero on error

int32_t onShapeChange(
nvinfer1::PluginTensorDesc const*,
int32_t,
nvinfer1::PluginTensorDesc const*,
int32_t
) noexcept override#

Called when input/output shapes change during runtime.

This is a no-op for this plugin type.

Parameters:

  • in – Input tensor descriptors (unused)

  • nbInputs – Number of inputs (unused)

  • out – Output tensor descriptors (unused)

  • nbOutputs – Number of outputs (unused)

Returns:

0 to signal success

nvinfer1::IPluginV3 *attachToContext(
nvinfer1::IPluginResourceContext*
) noexcept override#

Attach plugin to an execution context.

Parameters:

context – Plugin resource context (not used by this plugin)

Returns:

Cloned plugin attached to context

nvinfer1::PluginFieldCollection const *getFieldsToSerialize(
) noexcept override#

Get plugin fields for serialization.

Returns:

Field collection for serialization

void setPluginNamespace(char const *pluginNamespace) noexcept#

Set plugin namespace.

Parameters:

pluginNamespace – Namespace string

int32_t getNbTactics() noexcept override#

Number of custom tactics (valid cuteDSL variants for this N/K).

int32_t getValidTactics(
int32_t *tactics,
int32_t nbTactics
) noexcept override#

Fill the valid custom tactic ids (positive; = variant index + 1).

int32_t setTactic(int32_t tactic) noexcept override#

Select a tactic (>0 = cuteDSL variant id; 0 = default variant).

char const *getTimingCacheID() noexcept override#

Opt into TRT’s build timing cache so identical-shape int4 layers are autotuned once.

class Int4GroupwiseGemmPluginV2Creator : public nvinfer1::IPluginCreatorV3One#

Factory for creating Int4GroupwiseGemmPluginV2 instances.

Handles plugin registration and creation in TensorRT.

Public Functions

Int4GroupwiseGemmPluginV2Creator()#

Constructor.

~Int4GroupwiseGemmPluginV2Creator() override = default#

Destructor.

char const *getPluginName() const noexcept override#

Get plugin name.

Returns:

Plugin name string

char const *getPluginVersion() const noexcept override#

Get plugin version.

Returns:

Version string

nvinfer1::PluginFieldCollection const *getFieldNames(
) noexcept override#

Get plugin field names.

Returns:

Field collection

char const *getPluginNamespace() const noexcept override#

Get plugin namespace.

Returns:

Namespace string

void setPluginNamespace(char const *pluginNamespace) noexcept#

Set plugin namespace.

Parameters:

pluginNamespace – Namespace string

nvinfer1::IPluginV3 *createPlugin(
char const *name,
nvinfer1::PluginFieldCollection const *fc,
nvinfer1::TensorRTPhase
) noexcept override#

Create plugin from field collection.

Parameters:

  • name – Plugin name

  • fc – Field collection with parameters

  • phase – TensorRT phase (build or runtime) - not used by this plugin

Returns:

Created plugin instance