Building Custom TensorRT-LLM Plugins for Low-Latency FP8 Inference at Scale — code snippets

snippet-1.txt

#pragma once
#include <NvInferRuntime.h>
#include <vector>
#include <string>

namespace tensorrt_llm::plugins {

class FusedSwiGLUQuantPlugin : public nvinfer1::IPluginV2DynamicExt {
public:
    FusedSwiGLUQuantPlugin(float scale, bool has_dynamic_scale);
    FusedSwiGLUQuantPlugin(const void* data, size_t length);
    ~FusedSwiGLUQuantPlugin() override = default;

    // IPluginV2 methods
    const char* getPluginType() const noexcept override;
    const char* getPluginVersion() const noexcept override;
    int getNbOutputs() const noexcept override;
    nvinfer1::DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override;
    int initialize() noexcept override;
    void terminate() noexcept override;
    size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept override;
    int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
    
    // IPluginV2Ext methods
    nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept override;

    // IPluginV2DynamicExt methods
    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
    bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept override;
    void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept override;
    size_t getSerializationSize() const noexcept override;
    void serialize(void* buffer) const noexcept override;
    void setPluginNamespace(const char* pluginNamespace) noexcept override;
    const char* getPluginNamespace() const noexcept override;

private:
    float mScale;
    bool mHasDynamicScale;
    std::string mNamespace;
};

class FusedSwiGLUQuantPluginCreator : public nvinfer1::IPluginCreator {
public:
    FusedSwiGLUQuantPluginCreator();
    ~FusedSwiGLUQuantPluginCreator() override = default;
    const char* getPluginName() const noexcept override;
    const char* getPluginVersion() const noexcept override;
    const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override;
    nvinfer1::IPluginV2DynamicExt* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) noexcept override;
    nvinfer1::IPluginV2DynamicExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept override;
    void setPluginNamespace(const char* pluginNamespace) noexcept override;
    const char* getPluginNamespace() const noexcept override;

private:
    nvinfer1::PluginFieldCollection mFC;
    std::vector<nvinfer1::PluginField> mPluginAttributes;
    std::string mNamespace;
};

} // namespace tensorrt_llm::plugins

snippet-2.txt

#include <cuda_fp8.h>
#include <cuda_bf16.h>
#include <cuda_runtime.h>

// Vectorized SwiGLU + FP8 Quantization kernel
// Input:  input [N, 2 * H] in __nv_bfloat16
// Output: output [N, H] in __nv_fp8_e4m3
// scale:  Pointer to the FP32 quantization scale factor (y = x * scale)
__global__ void fused_swiglu_quant_kernel(
    const __nv_bfloat16* __restrict__ input,
    __nv_fp8_e4m3* __restrict__ output,
    const float* __restrict__ scale,
    const int N,
    const int H) {

    // Map CUDA block variables to rows (tokens)
    const int row = blockIdx.y;
    if (row >= N) return;

    const __nv_bfloat16* row_in = input + row * 2 * H;
    __nv_fp8_e4m3* row_out = output + row * H;

    const float s = *scale;

    // Each thread processes 8 elements of the output H (meaning 16 elements of input 2*H)
    for (int col = (blockIdx.x * blockDim.x + threadIdx.x) * 8; col < H; col += blockDim.x * gridDim.x * 8) {
        
        // Vectorized 128-bit load for the first half of the input (gate inputs)
        const float4* gate_in_ptr = reinterpret_cast<const float4*>(row_in + col);
        float4 gate_val_raw = *gate_in_ptr;
        const __nv_bfloat16* gate_h = reinterpret_cast<const __nv_bfloat16*>(&gate_val_raw);

        // Vectorized 128-bit load for the second half of the input (value inputs)
        const float4* val_in_ptr = reinterpret_cast<const float4*>(row_in + H + col);
        float4 val_val_raw = *val_in_ptr;
        const __nv_bfloat16* val_h = reinterpret_cast<const __nv_bfloat16*>(&val_val_raw);

        __nv_fp8_e4m3 out_fp8[8];

        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            float g = __bfloat162float(gate_h[i]);
            float v = __bfloat162float(val_h[i]);

            // Swish activation: swish(g) = g * sigmoid(g) = g / (1 + exp(-g))
            float swish = g / (1.0f + __expf(-g));
            float swiglu = swish * v;

            // Apply FP8 scaling
            float scaled_val = swiglu * s;

            // Saturate to FP8 E4M3 dynamic range: max representable value is 448.0f
            scaled_val = __fmaxf(-448.0f, __fminf(scaled_val, 448.0f));
            
            out_fp8[i] = __nv_fp8_e4m3(scaled_val);
        }

        // Vectorized 64-bit store of 8 FP8 elements to avoid memory alignment stalls
        uint2* out_ptr = reinterpret_cast<uint2*>(row_out + col);
        *out_ptr = *reinterpret_cast<uint2*>(out_fp8);
    }
}

snippet-3.txt

#include <cuda_runtime.h>
#include <cuda_fp8.h>
#include <cuda_bf16.h>

void launch_fused_swiglu_quant(
    const __nv_bfloat16* input,
    __nv_fp8_e4m3* output,
    const float* scale,
    int N,
    int H,
    cudaStream_t stream) {

    // Since we process 8 output elements per thread, thread block size governs output slice width.
    // For H = 4096, 256 threads will cover 2048 elements per loop.
    int threads_per_block = 256;
    if (H / 8 < threads_per_block) {
        threads_per_block = (H / 8 + 31) / 32 * 32; // Round to warp alignment
        if (threads_per_block == 0) threads_per_block = 32;
    }

    // Grid dimension X covers the hidden dimension; Grid Y processes the tokens (rows)
    int blocks_x = (H + (threads_per_block * 8) - 1) / (threads_per_block * 8);
    dim3 grid(blocks_x, N);
    dim3 block(threads_per_block);

    fused_swiglu_quant_kernel<<<grid, block, 0, stream>>>(
        input,
        output,
        scale,
        N,
        H
    );
}

snippet-4.txt

#include "fp8_activation_plugin.h"
#include <cassert>

namespace tensorrt_llm::plugins {

FusedSwiGLUQuantPlugin::FusedSwiGLUQuantPlugin(float scale, bool has_dynamic_scale)
    : mScale(scale), mHasDynamicScale(has_dynamic_scale) {}

const char* FusedSwiGLUQuantPlugin::getPluginType() const noexcept {
    return "FusedSwiGLUQuant";
}

const char* FusedSwiGLUQuantPlugin::getPluginVersion() const noexcept {
    return "1";
}

int FusedSwiGLUQuantPlugin::getNbOutputs() const noexcept {
    return 1;
}

nvinfer1::DimsExprs FusedSwiGLUQuantPlugin::getOutputDimensions(
    int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept {
    
    assert(nbInputs == 2);
    assert(outputIndex == 0);

    nvinfer1::DimsExprs outputDims = inputs[0];
    
    // The input dimension is [N, 2 * H]. The output dimension must be [N, H].
    // We divide the last dimension of the input tensor by 2.
    auto two = exprBuilder.constant(2);
    outputDims.d[outputDims.nbDims - 1] = exprBuilder.operation(
        nvinfer1::DimensionOperation::kFLOOR_DIV, 
        *inputs[0].d[inputs[0].nbDims - 1], 
        *two
    );
    return outputDims;
}

bool FusedSwiGLUQuantPlugin::supportsFormatCombination(
    int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept {
    
    assert(nbInputs == 2);
    assert(nbOutputs == 1);

    if (pos == 0) { // Input tensor
        return (inOut[pos].type == nvinfer1::DataType::kHALF || inOut[pos].type == nvinfer1::DataType::kBF16)
            && inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
    }
    else if (pos == 1) { // Scale factor tensor
        return inOut[pos].type == nvinfer1::DataType::kFLOAT
            && inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
    }
    else if (pos == 2) { // Output tensor
        return inOut[pos].type == nvinfer1::DataType::kFP8
            && inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
    }
    return false;
}

void FusedSwiGLUQuantPlugin::configurePlugin(
    const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs,
    const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept {
    
    assert(nbInputs == 2);
    assert(nbOutputs == 1);
    
    const auto& inputDims = in[0].desc.dims;
    const auto& outputDims = out[0].desc.dims;
    
    int lastDimInput = inputDims.d[inputDims.nbDims - 1];
    int lastDimOutput = outputDims.d[outputDims.nbDims - 1];
    assert(lastDimInput == 2 * lastDimOutput);
}

size_t FusedSwiGLUQuantPlugin::getWorkspaceSize(
    const nvinfer1::PluginTensorDesc* inputs, int nbInputs,
    const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept {
    // Fused operations are computed in registers, requiring 0 bytes of temporary HBM workspace.
    return 0;
}

nvinfer1::DataType FusedSwiGLUQuantPlugin::getOutputDataType(
    int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept {
    assert(index == 0);
    return nvinfer1::DataType::kFP8;
}

nvinfer1::IPluginV2DynamicExt* FusedSwiGLUQuantPlugin::clone() const noexcept {
    auto* plugin = new FusedSwiGLUQuantPlugin(mScale, mHasDynamicScale);
    plugin->setPluginNamespace(mNamespace.c_str());
    return plugin;
}

int FusedSwiGLUQuantPlugin::initialize() noexcept {
    return 0;
}

void FusedSwiGLUQuantPlugin::terminate() noexcept {}

void FusedSwiGLUQuantPlugin::setPluginNamespace(const char* pluginNamespace) noexcept {
    mNamespace = pluginNamespace;
}

const char* FusedSwiGLUQuantPlugin::getPluginNamespace() const noexcept {
    return mNamespace.c_str();
}

} // namespace tensorrt_llm::plugins

snippet-5.txt

#include "fp8_activation_plugin.h"

// Forward declaration of the CUDA launcher
void launch_fused_swiglu_quant(
    const __nv_bfloat16* input,
    __nv_fp8_e4m3* output,
    const float* scale,
    int N,
    int H,
    cudaStream_t stream);

namespace tensorrt_llm::plugins {

int FusedSwiGLUQuantPlugin::enqueue(
    const nvinfer1::PluginTensorDesc* inputDesc,
    const nvinfer1::PluginTensorDesc* outputDesc,
    const void* const* inputs,
    void* const* outputs,
    void* workspace,
    cudaStream_t stream) noexcept {
    
    // Determine the number of tokens (N) by flattening dimensions [0, nbDims - 2]
    int N = 1;
    for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) {
        N *= inputDesc[0].dims.d[i];
    }
    // Hidden size (H) is half the input hidden size
    int H = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1] / 2;

    const __nv_bfloat16* input = static_cast<const __nv_bfloat16*>(inputs[0]);
    const float* scale = static_cast<const float*>(inputs[1]);
    __nv_fp8_e4m3* output = static_cast<__nv_fp8_e4m3*>(outputs[0]);

    launch_fused_swiglu_quant(input, output, scale, N, H, stream);

    return 0;
}

} // namespace tensorrt_llm::plugins

snippet-6.txt

#include "fp8_activation_plugin.h"
#include <cstring>

namespace tensorrt_llm::plugins {

size_t FusedSwiGLUQuantPlugin::getSerializationSize() const noexcept {
    return sizeof(mScale) + sizeof(mHasDynamicScale);
}

void FusedSwiGLUQuantPlugin::serialize(void* buffer) const noexcept {
    char* d = static_cast<char*>(buffer);
    std::memcpy(d, &mScale, sizeof(mScale));
    d += sizeof(mScale);
    std::memcpy(d, &mHasDynamicScale, sizeof(mHasDynamicScale));
}

// Deserialization Constructor
FusedSwiGLUQuantPlugin::FusedSwiGLUQuantPlugin(const void* data, size_t length) {
    const char* d = static_cast<const char*>(data);
    std::memcpy(&mScale, d, sizeof(mScale));
    d += sizeof(mScale);
    std::memcpy(&mHasDynamicScale, d, sizeof(mHasDynamicScale));
}

FusedSwiGLUQuantPluginCreator::FusedSwiGLUQuantPluginCreator() {
    mPluginAttributes.clear();
    mPluginAttributes.emplace_back(nvinfer1::PluginField("scale", nullptr, nvinfer1::PluginFieldType::kFLOAT32, 1));
    mPluginAttributes.emplace_back(nvinfer1::PluginField("has_dynamic_scale", nullptr, nvinfer1::PluginFieldType::kINT32, 1));
    
    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

const char* FusedSwiGLUQuantPluginCreator::getPluginName() const noexcept {
    return "FusedSwiGLUQuant";
}

const char* FusedSwiGLUQuantPluginCreator::getPluginVersion() const noexcept {
    return "1";
}

const nvinfer1::PluginFieldCollection* FusedSwiGLUQuantPluginCreator::getFieldNames() noexcept {
    return &mFC;
}

nvinfer1::IPluginV2DynamicExt* FusedSwiGLUQuantPluginCreator::createPlugin(
    const char* name, const nvinfer1::PluginFieldCollection* fc) noexcept {
    
    float scale = 1.0f;
    bool hasDynamicScale = false;

    for (int i = 0; i < fc->nbFields; ++i) {
        std::string fieldName(fc->fields[i].name);
        if (fieldName == "scale") {
            scale = *static_cast<const float*>(fc->fields[i].data);
        } else if (fieldName == "has_dynamic_scale") {
            hasDynamicScale = *static_cast<const bool*>(fc->fields[i].data);
        }
    }
    
    auto* plugin = new FusedSwiGLUQuantPlugin(scale, hasDynamicScale);
    plugin->setPluginNamespace(mNamespace.c_str());
    return plugin;
}

nvinfer1::IPluginV2DynamicExt* FusedSwiGLUQuantPluginCreator::deserializePlugin(
    const char* name, const void* serialData, size_t serialLength) noexcept {
    
    auto* plugin = new FusedSwiGLUQuantPlugin(serialData, serialLength);
    plugin->setPluginNamespace(mNamespace.c_str());
    return plugin;
}

void FusedSwiGLUQuantPluginCreator::setPluginNamespace(const char* pluginNamespace) noexcept {
    mNamespace = pluginNamespace;
}

const char* FusedSwiGLUQuantPluginCreator::getPluginNamespace() const noexcept {
    return mNamespace.c_str();
}

// Register the creator with the global TensorRT plugin registry
REGISTER_TENSORRT_PLUGIN(FusedSwiGLUQuantPluginCreator);

} // namespace tensorrt_llm::plugins

snippet-8.txt

// Example of checking memory alignment before running vectorized loads
bool is_aligned = (reinterpret_cast<uintptr_t>(input) % 16 == 0) && (H % 8 == 0);
if (is_aligned) {
    // Launch vectorized kernel
} else {
    // Launch slow fallback scalar kernel
}