Low Rank Candidate Head Plugin

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

Low-rank candidate head of an EAGLE candidate draft: inputs hidden [R, H], down_proj [r, H], up_proj [V, r] and lm_head [V, H] (fp16); outputs support_ids [R, Kc] int32 and support_log_probs [R, Kc] fp32.

Only num_candidates (Kc) is an ONNX attribute. configurePlugin resolves hidden_size, rank, vocab_size and max_rows from the weight dims and the largest profile maximum of R it is configured for, plus workspace_bytes for the build device; all six are serialized. enqueue runs the head in the current device’s preferred configuration or, when the engine reserved less workspace, in a smaller one that fits (results do not depend on it); the first enqueue on a device picks it, and every enqueue fails without launching if none fits. attachToContext only clones: TensorRT attaches a plugin layer that the compiler backend does not absorb before configuring it.

Public Functions

LowRankCandidateHeadPlugin(

std::string const &name,

int32_t numCandidates,

int32_t hiddenSize,

int32_t rank,

int32_t vocabSize,

int32_t maxRows,

int64_t workspaceBytes

)

LowRankCandidateHeadPlugin(LowRankCandidateHeadPlugin const &other)

~LowRankCandidateHeadPlugin() 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

void setPluginNamespace(char const *pluginNamespace) noexcept

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 *inputs,

int32_t nbInputs,

nvinfer1::DynamicPluginTensorDesc const *outputs,

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 *inputs,

int32_t nbInputs,

nvinfer1::PluginTensorDesc const *outputs,

int32_t nbOutputs

) noexcept override

nvinfer1::IPluginV3 *attachToContext(

nvinfer1::IPluginResourceContext *context

) noexcept override

nvinfer1::PluginFieldCollection const *getFieldsToSerialize(

) noexcept override

class LowRankCandidateHeadPluginCreator : public nvinfer1::IPluginCreatorV3One

Public Functions

LowRankCandidateHeadPluginCreator()

~LowRankCandidateHeadPluginCreator() 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 *fields,

nvinfer1::TensorRTPhase phase

) noexcept override

The build phase needs only num_candidates; the runtime phase needs every serialized field. Other fields (such as the _outputs attribute that onnxscript leaves on the node) are ignored.