chore: import upstream snapshot with attribution

2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -0,0 +1,81 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+cmake_minimum_required(VERSION 3.31 FATAL_ERROR)
+project(CustomHardMax LANGUAGES CXX CUDA)
+
+find_package(CUDAToolkit REQUIRED)
+
+if(NOT MSVC)
+    # Enable all compile warnings
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wno-long-long -pedantic -Wno-deprecated-declarations")
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wno-deprecated-declarations")
+endif()
+
+# Sets variable to a value if variable is unset.
+macro(set_ifndef var val)
+    if(NOT ${var})
+        set(${var} ${val})
+    endif()
+    message(STATUS "Configurable variable ${var} set to ${${var}}")
+endmacro()
+
+# -------- CONFIGURATION --------
+if(NOT MSVC)
+    set_ifndef(TRT_LIB /usr/lib/x86_64-linux-gnu)
+    set_ifndef(TRT_INCLUDE /usr/include/x86_64-linux-gnu)
+
+endif()
+
+# Find dependencies:
+message("\nThe following variables are derived from the values of the previous variables unless provided explicitly:\n")
+
+# TensorRT’s nvinfer lib
+find_library(
+    _NVINFER_LIB nvinfer
+    HINTS ${TRT_LIB}
+    PATH_SUFFIXES lib lib64)
+set_ifndef(NVINFER_LIB ${_NVINFER_LIB})
+
+# -------- BUILDING --------
+
+add_definitions(-DTENSORRT_BUILD_LIB)
+
+# Add include directories
+file(REAL_PATH ${CMAKE_SOURCE_DIR} CMAKE_SOURCE_DIR_REALPATH)
+
+get_filename_component(SAMPLES_COMMON_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../common/ ABSOLUTE)
+get_filename_component(SHARED_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../../shared ABSOLUTE)
+get_filename_component(SAMPLES_DIR ${CMAKE_SOURCE_DIR_REALPATH}/../../ ABSOLUTE)
+# Define Hardmax plugin library target
+add_library(
+    customHardmaxPlugin MODULE
+    ${SAMPLES_COMMON_DIR}/logger.cpp ${SHARED_DIR}/utils/fileLock.cpp
+    ${CMAKE_SOURCE_DIR_REALPATH}/plugin/customHardmaxPlugin.cpp ${CMAKE_SOURCE_DIR_REALPATH}/plugin/customHardmaxPlugin.h)
+
+target_include_directories(customHardmaxPlugin PRIVATE
+    ${CUDAToolkit_INCLUDE_DIRS} ${TRT_INCLUDE} ${CMAKE_SOURCE_DIR_REALPATH}/plugin/
+    ${SAMPLES_COMMON_DIR} ${SAMPLES_DIR} ${SHARED_DIR})
+
+# Use C++11
+target_compile_features(customHardmaxPlugin PUBLIC cxx_std_17)
+
+# Link TensorRT’s nvinfer lib
+target_link_libraries(customHardmaxPlugin PRIVATE ${NVINFER_LIB})
+target_link_libraries(customHardmaxPlugin PRIVATE CUDA::cudart_static)
+target_link_libraries(customHardmaxPlugin PRIVATE CUDA::cublas)
+target_link_libraries(customHardmaxPlugin PRIVATE CUDA::cuda_driver)
@@ -0,0 +1,177 @@
+# Adding A Custom Layer Implementation to Your ONNX Network
+
+**Table Of Contents**
+- [Description](#description)
+- [How does this sample work?](#how-does-this-sample-work)
+- [Prerequisites](#prerequisites)
+- [Download and preprocess the ONNX model](#download-the-onnx-model)
+- [Running the sample](#running-the-sample)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+This sample, `onnx_custom_plugin`, demonstrates how to use plugins written in C++ with the TensorRT Python bindings and ONNX Parser. This sample uses the [BiDAF Model](https://github.com/onnx/models/tree/main/text/machine_comprehension/bidirectional_attention_flow) from ONNX Model Zoo.
+
+## How does this sample work?
+
+This sample implements a Hardmax layer using cuBLAS, wraps the implementation in a TensorRT plugin (with a corresponding plugin creator) and then generates a shared library module containing its code. The user then dynamically loads this library in Python, which causes the plugin to be registered in TensorRT's PluginRegistry and makes it available to the ONNX parser.
+
+This sample includes:
+
+`plugin/`
+This directory contains files for the Hardmax layer plugin.
+
+`customHardmaxPlugin.cpp`
+A custom TensorRT plugin implementation.
+
+`customHardmaxPlugin.h`
+The Hardmax Plugin headers.
+
+`model.py`
+This script downloads the BiDAF onnx model and uses Onnx Graphsurgeon to replace layers unsupported by TensorRT.
+
+`sample.py`
+This script loads the ONNX model and performs inference using TensorRT.
+
+`load_plugin_lib.py`
+This script contains a helper function to load the customHardmaxPlugin library in Python.
+
+`test_custom_hardmax_plugin.py`
+This script tests the Hardmax Plugin against a reference numpy implementation.
+
+`requirements.txt`
+This file specifies all the Python packages required to run this Python sample.
+
+## Prerequisites
+
+For specific software versions, see the [TensorRT Installation Guide](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html).
+
+1. Install the dependencies for Python.
+
+```bash
+pip3 install -r requirements.txt
+```
+
+2. [Install CMake](https://cmake.org/download/).
+
+3. [Install Cublas](https://developer.nvidia.com/cublas).
+
+4. (For Windows builds) [Visual Studio](https://visualstudio.microsoft.com/vs/older-downloads/) 2017 Community or Enterprise edition
+
+## Download and preprocess the ONNX model
+
+Run the model script to download the BiDAF model from the Onnx Model Zoo. The script will replace the `Hardmax` layer with an op called `CustomHardmax` to match the custom Plugin name. It will also replace the unsupported `Compress` node with an equivalent operation, and remove the `CategoryMapper` nodes which do a String-to-Int conversion of the model inputs.
+
+```bash
+python3 model.py
+```
+
+## Running the sample
+
+1.  Build the plugin and its corresponding Python bindings.
+
+   - On Linux, run:
+      ```bash
+      mkdir build && pushd build
+      cmake .. && make -j
+      popd
+      ```
+
+      **NOTE:** If any of the dependencies are not installed in their default locations, you can manually specify them. For example:
+      ```bash
+      cmake .. -DCMAKE_CUDA_COMPILER=/usr/local/cuda-x.x/bin/nvcc # (Or adding /path/to/nvcc into $PATH)
+               -DCUDA_INC_DIR=/usr/local/cuda-x.x/include/  # (Or adding /path/to/cuda/include into $CPLUS_INCLUDE_PATH)
+               -DTRT_LIB=/path/to/tensorrt/lib/
+               -DTRT_INCLUDE=/path/to/tensorrt/include/
+      ```
+
+   - On Windows, run the following in Powershell, replacing paths appropriately:
+      ```ps1
+      mkdir build; pushd build
+      cmake .. -G "Visual Studio 15 Win64" /
+         -DTRT_LIB=C:\path\to\tensorrt\lib /
+         -DTRT_INCLUDE=C:\path\to\tensorrt\lib /
+         -DCUDA_INC_DIR="C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v<CUDA_VERSION>\include" /
+         -DCUDA_LIB_DIR="C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v<CUDA_VERSION>\lib\x64"
+      # NOTE: msbuild is usually located under C:\Program Files (x86)\Microsoft Visual Studio\2017\<EDITION>\MSBuild\<VERSION>\Bin
+      #   You should add this path to your PATH environment variable.
+      msbuild ALL_BUILD.vcxproj
+      popd
+      ```
+
+   The command `cmake ..` displays a complete list of configurable variables. If a variable is set to `VARIABLE_NAME-NOTFOUND`, then you’ll need to specify it manually or set the variable it is derived from correctly.
+
+2.  Run inference using TensorRT with the custom Hardmax plugin implementation:
+   ```bash
+   python3 sample.py
+   ```
+
+3.  Verify that the sample ran successfully.
+   ```
+   === Testing ===
+
+   Input context: Garry the lion is 5 years old. He lives in the savanna.
+   Input query: Where does the lion live?
+   Model prediction:  savanna
+
+   Input context: A quick brown fox jumps over the lazy dog.
+   Input query: What color is the fox?
+   Model prediction:  brown
+   ```
+
+   The model can also be run interactively:
+   ```bash
+   python3 sample.py --interactive
+   ```
+
+   The context and query can then be entered from the command line:
+
+   ```
+   === Testing ===
+   Enter context: Waldo wears a striped shirt. He also wears glasses.
+   Enter query: Who wears glasses?
+   Model prediction:  waldo
+   ```
+
+# Additional resources
+
+The following resources provide a deeper understanding about getting started with TensorRT using Python:
+
+**Model**
+- [BiDAF model](https://allenai.github.io/bi-att-flow/)
+
+**Documentation**
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+
+# License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+# Changelog
+March 2026
+  - Migrate HardmaxPlugin from IPluginV2DynamicExt to IPluginV3.
+
+October 2025
+  - Migrate to strongly typed APIs.
+
+August 2025:
+  - Removed support for Python versions < 3.10.
+
+January 2024:
+  - Create cublas handle with cublasCreate instead of using the cublasContext argument from attachToContext.
+  - Added the Cublas library as a prerequisite.
+
+August 2023:
+  - Update ONNX version support to 1.14.0
+  - Removed support for Python versions < 3.8.
+
+September 2022: This `README.md` file was created and reviewed.
+
+# Known issues
+
+There are no known issues in this sample.
@@ -0,0 +1,58 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import ctypes
+
+WORKING_DIR = os.environ.get("TRT_WORKING_DIR") or os.path.dirname(
+    os.path.realpath(__file__)
+)
+IS_WINDOWS = os.name == "nt"
+if IS_WINDOWS:
+    HARDMAX_PLUGIN_LIBRARY_NAME = "customHardmaxPlugin.dll"
+    HARDMAX_PLUGIN_LIBRARY = [
+        os.path.join(WORKING_DIR, "build", "Debug", HARDMAX_PLUGIN_LIBRARY_NAME),
+        os.path.join(WORKING_DIR, "build", "Release", HARDMAX_PLUGIN_LIBRARY_NAME),
+    ]
+else:
+    HARDMAX_PLUGIN_LIBRARY_NAME = "libcustomHardmaxPlugin.so"
+    HARDMAX_PLUGIN_LIBRARY = [
+        os.path.join(WORKING_DIR, "build", HARDMAX_PLUGIN_LIBRARY_NAME)
+    ]
+
+
+def load_plugin_lib():
+    for plugin_lib in HARDMAX_PLUGIN_LIBRARY:
+        if os.path.isfile(plugin_lib):
+            try:
+                # Python specifies that winmode is 0 by default, but some implementations
+                # incorrectly default to None instead. See:
+                # https://docs.python.org/3.8/library/ctypes.html
+                # https://github.com/python/cpython/blob/3.10/Lib/ctypes/__init__.py#L343
+                ctypes.CDLL(plugin_lib, winmode=0)
+            except TypeError:
+                # winmode only introduced in python 3.8
+                ctypes.CDLL(plugin_lib)
+            return
+
+    raise IOError(
+        "\n{}\n{}\n{}\n".format(
+            "Failed to load library ({}).".format(HARDMAX_PLUGIN_LIBRARY_NAME),
+            "Please build the Hardmax sample plugin.",
+            "For more information, see the included README.md",
+        )
+    )
@@ -0,0 +1,129 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+
+import wget
+import onnx
+import onnx_graphsurgeon as gs
+
+MODEL_URL = "https://github.com/onnx/models/raw/e77240a62df68ed13e3138a5812553a552b857bb/text/machine_comprehension/bidirectional_attention_flow/model/bidaf-9.onnx"
+
+WORKING_DIR = os.environ.get("TRT_WORKING_DIR") or os.path.dirname(
+    os.path.realpath(__file__)
+)
+MODEL_DIR = os.path.join(WORKING_DIR, "models")
+RAW_MODEL_PATH = os.path.join(MODEL_DIR, "bidaf-9.onnx")
+TRT_MODEL_PATH = os.path.join(MODEL_DIR, "bidaf-9-trt.onnx")
+
+
+def _do_graph_surgery(raw_model_path, trt_model_path):
+    graph = gs.import_onnx(onnx.load(raw_model_path))
+
+    # Replace unsupported Hardmax with our CustomHardmax op
+    hardmax_node = None
+    for node in graph.nodes:
+        if node.op == "Hardmax":
+            node.op = "CustomHardmax"
+            hardmax_node = node
+    assert hardmax_node is not None, "Model does not contain a Hardmax node"
+
+    # The original onnx model also uses another unsupported op called "Compress".
+    # "Compress" returns values from the first tensor for all indices which evaluate to
+    # True in the second tensor. In our case the second Tensor is the output of Hardmax,
+    # so exactly one index will evaluate to true because the value at it will be 1, and
+    # all other values will be 0. We can achieve the same result as "Compress" by taking the
+    # dot product of our value tensor and the Hardmax output.
+    #
+    # So, we will replace the subgraph Compress(Transpose_29, Cast(Reshape(Hardmax)))
+    # with the subgraph Einsum(Transpose_29, Hardmax) where the equation in Einsum takes the dot product.
+    node_by_name = {node.name: node for node in graph.nodes}
+    transpose_node = node_by_name["Transpose_29"]
+    compress_node = node_by_name["Compress_31"]
+
+    einsum_node = gs.Node(
+        "Einsum",
+        "Dot_of_Hardmax_and_Transpose",
+        attrs={"equation": "ij,ij->i"},  # "Dot product" of 2d tensors
+        inputs=[hardmax_node.outputs[0], transpose_node.outputs[0]],
+        outputs=[compress_node.outputs[0]],
+    )
+    graph.nodes.append(einsum_node)
+
+    # Separate the old subgraph which will be deleted with graph.cleanup()
+    hardmax_node.o().inputs.clear()
+    transpose_node.o().inputs.clear()
+    compress_node.outputs.clear()
+
+    # Also remove the CategoryMapper nodes which convert strings to integers as the first step in the model.
+    # We need to convert the following structure:
+    #
+    #      Input as                        Converted to
+    #   String tokens                     Integer tokens
+    #  ---------------->[CategoryMapper]------------------>[Rest of Model]
+    #
+    # into the following:
+    #
+    #      Input as
+    #   Integer tokens
+    #  ------------------>[Rest of Model]
+    #
+    # Later we will feed the model the integer tokens directly.
+    # Note: list conversion is necessary because we modify graph.nodes in the for loop.
+    category_mapper_nodes = [
+        node for node in graph.nodes if node.op == "CategoryMapper"
+    ]
+    for node in category_mapper_nodes:
+        # Remove CategoryMapper node from onnx graph
+        graph.nodes.remove(node)
+
+        # Also remove references its inputs in the graph's inputs
+        for input_tensor in node.inputs:
+            graph.inputs.remove(input_tensor)
+
+        # The graph's new inputs are the Integer tokens output by CategoryMapper
+        graph.inputs += node.outputs
+
+        # Save String->Int map
+        with open(node.name + ".json", "w") as fp:
+            json.dump(node.attrs, fp)
+
+    graph.cleanup().toposort()
+    onnx.save(gs.export_onnx(graph), trt_model_path)
+
+
+def make_trt_compatible_onnx_model():
+    os.makedirs(MODEL_DIR, exist_ok=True)
+    if not os.path.exists(RAW_MODEL_PATH):
+        wget.download(MODEL_URL, out=RAW_MODEL_PATH)
+        print("\nDownloaded BiDAF model from Onnx Model Zoo")
+    print("Performing graph surgery on Onnx Model Zoo BiDAF model")
+    _do_graph_surgery(RAW_MODEL_PATH, TRT_MODEL_PATH)
+    print("Graph Surgery complete!")
+
+
+def main():
+    if os.path.exists(TRT_MODEL_PATH):
+        print("TRT-compatible onnx model already exists!")
+    else:
+        print("TRT-compatible onnx model not found, generating...")
+        make_trt_compatible_onnx_model()
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,394 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "customHardmaxPlugin.h"
+#include "NvInferPlugin.h"
+#include "common.h" // volume(), ASSERT
+#include "logger.h" // sample::gLogError
+#include <cuda.h>
+#include <memory>
+#include <string_view>
+
+using namespace nvinfer1;
+
+#define CUDRIVER_CALL(call)                                                                                            \
+    {                                                                                                                  \
+        cudaError_enum s_ = call;                                                                                      \
+        if (s_ != CUDA_SUCCESS)                                                                                        \
+        {                                                                                                              \
+            char const *errName_, *errDesc_;                                                                           \
+            cuGetErrorName(s_, &errName_);                                                                             \
+            cuGetErrorString(s_, &errDesc_);                                                                           \
+            sample::gLogError << "CUDA Error: " << errName_ << " " << errDesc_ << std::endl;                           \
+            return s_;                                                                                                 \
+        }                                                                                                              \
+    }
+
+#define CUDA_CALL(call)                                                                                                \
+    {                                                                                                                  \
+        cudaError_t s_ = call;                                                                                         \
+        if (s_ != cudaSuccess)                                                                                         \
+        {                                                                                                              \
+            sample::gLogError << "CUDA Error: " << cudaGetErrorName(s_) << " " << cudaGetErrorString(s_) << std::endl; \
+            return s_;                                                                                                 \
+        }                                                                                                              \
+    }
+
+#define CUBLAS_CALL(call)                                                                                              \
+    {                                                                                                                  \
+        cublasStatus_t s_ = call;                                                                                      \
+        if (s_ != CUBLAS_STATUS_SUCCESS)                                                                               \
+        {                                                                                                              \
+            sample::gLogError << "cuBLAS Error: " << s_ << std::endl;                                                  \
+            return s_;                                                                                                 \
+        }                                                                                                              \
+    }
+
+REGISTER_TENSORRT_PLUGIN(HardmaxPluginCreator);
+
+namespace
+{
+constexpr char const* kHARDMAX_NAME{"CustomHardmax"};
+constexpr char const* kHARDMAX_VERSION{"1"};
+} // namespace
+
+HardmaxPlugin::HardmaxPlugin(int32_t axis)
+    : mAxis(axis)
+{
+}
+
+HardmaxPlugin::HardmaxPlugin(HardmaxPlugin const& other)
+    : mNamespace(other.mNamespace)
+    , mAxisSize(other.mAxisSize)
+    , mDimProductOuter(other.mDimProductOuter)
+    , mDimProductInner(other.mDimProductInner)
+    , mCublas(nullptr)
+    , mAxis(other.mAxis)
+{
+}
+
+HardmaxPlugin::~HardmaxPlugin()
+{
+    if (mCublas)
+    {
+        cublasDestroy(mCublas);
+    }
+}
+
+// IPluginV3 methods
+
+IPluginCapability* HardmaxPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept
+{
+    if (type == PluginCapabilityType::kBUILD)
+    {
+        return static_cast<IPluginV3OneBuild*>(this);
+    }
+    if (type == PluginCapabilityType::kRUNTIME)
+    {
+        return static_cast<IPluginV3OneRuntime*>(this);
+    }
+    ASSERT(type == PluginCapabilityType::kCORE);
+    return static_cast<IPluginV3OneCore*>(this);
+}
+
+IPluginV3* HardmaxPlugin::clone() noexcept
+{
+    auto plugin = std::make_unique<HardmaxPlugin>(*this);
+    return plugin.release();
+}
+
+// IPluginV3OneCore methods
+
+char const* HardmaxPlugin::getPluginName() const noexcept
+{
+    return kHARDMAX_NAME;
+}
+
+char const* HardmaxPlugin::getPluginVersion() const noexcept
+{
+    return kHARDMAX_VERSION;
+}
+
+char const* HardmaxPlugin::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+// IPluginV3OneBuild methods
+
+int32_t HardmaxPlugin::getNbOutputs() const noexcept
+{
+    return 1;
+}
+
+int32_t HardmaxPlugin::getOutputDataTypes(
+    DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
+{
+    ASSERT(inputTypes != nullptr);
+    ASSERT(nbInputs == 1);
+    ASSERT(nbOutputs == 1);
+    outputTypes[0] = inputTypes[0];
+    return 0;
+}
+
+int32_t HardmaxPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
+    int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept
+{
+    ASSERT(nbInputs == 1);
+    ASSERT(nbOutputs == 1);
+    outputs[0] = inputs[0];
+    return 0;
+}
+
+bool HardmaxPlugin::supportsFormatCombination(
+    int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
+{
+    ASSERT(inOut && pos < (nbInputs + nbOutputs));
+
+    // Type changes are not allowed
+    if (inOut[0].desc.type != inOut[pos].desc.type)
+    {
+        return false;
+    }
+
+    return inOut[pos].desc.type == DataType::kFLOAT && inOut[pos].desc.format == PluginFormat::kLINEAR;
+}
+
+int32_t HardmaxPlugin::configurePlugin(
+    DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    ASSERT(nbInputs == 1);
+    ASSERT(nbOutputs == 1);
+
+    Dims const& inDims = in[0].desc.dims;
+
+    // Normalize negative axis to positive
+    if (mAxis < 0)
+    {
+        mAxis += inDims.nbDims;
+        ASSERT(mAxis >= 0);
+    }
+    ASSERT(inDims.nbDims > mAxis);
+
+    return 0;
+}
+
+size_t HardmaxPlugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+    DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
+{
+    ASSERT(mAxis >= 0);
+    // Two arrays are needed:
+    // 1. For the contents of the working axis
+    // 2. For an array of 1's
+    return 2 * inputs[0].max.d[mAxis] * sizeof(float);
+}
+
+// IPluginV3OneRuntime methods
+
+int32_t HardmaxPlugin::onShapeChange(
+    PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    ASSERT(nbInputs == 1);
+    ASSERT(nbOutputs == 1);
+
+    Dims const& inDims = in[0].dims;
+
+    // Axis should already be normalized by configurePlugin, but handle it regardless to be safe.
+    if (mAxis < 0)
+    {
+        mAxis += inDims.nbDims;
+        ASSERT(mAxis >= 0);
+    }
+    ASSERT(inDims.nbDims > mAxis);
+
+    mDimProductOuter = samplesCommon::volume(inDims, 0, mAxis);
+    mAxisSize = inDims.d[mAxis];
+    mDimProductInner = samplesCommon::volume(inDims, mAxis + 1, inDims.nbDims);
+
+    return 0;
+}
+
+int32_t HardmaxPlugin::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
+    void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
+{
+    if (inputDesc[0].type != DataType::kFLOAT)
+    {
+        return -1;
+    }
+
+    CUBLAS_CALL(cublasSetStream(mCublas, stream));
+
+    auto const* data = static_cast<float const*>(inputs[0]);
+    auto* result = static_cast<float*>(outputs[0]);
+
+    // Make sure output is initialized to all 0's.
+    // Later we will set the correct outputs to be 1's and not touch the rest.
+    CUDA_CALL(cudaMemsetAsync(result, 0, mDimProductOuter * mDimProductInner * mAxisSize * sizeof(float), stream));
+
+    // We use the workspace in the case that the first call to 'cublasIsamax' is insufficient.
+    // The first half of the workspace we use to copy the values of the axis into, so that we can
+    // subtract out the minimum value and call 'cublasIsamax' again. See the comment below.
+    // The second half of the workspace will be a costant array of 1's, necessary for our cublasSaxpy call.
+    auto* const axisFlat = static_cast<float* const>(workspace);
+    float* const ones = axisFlat + mAxisSize;
+    float const one = 1.0F;
+    CUDRIVER_CALL(cuMemsetD32Async(CUdeviceptr(ones), *reinterpret_cast<int const*>(&one), mAxisSize, stream));
+
+    // This plugin works by parallelizing the argmax operation along a single axis.
+    // This is efficient when the axis size is very large compared to the other dimensions.
+    //
+    // Consider an input shape (1, 512, 3) with axis = 1. This plugin will perform well because
+    // the work which is parallelized is over the large 512-element-long axis, and the work that is done
+    // serially is over the small 1-element-long and 3-element-long axes.
+    //
+    // However, when the axis size is small compared to the other dimensions, this plugin will be very
+    // inefficient. If the input shape is (1, 512, 3) and the hardmax is over axis = 2, then
+    // the work is parallelized over the small 3-element-long axis and the work is done serially over
+    // the large 512-element-long axis. A smarter plugin would try to recognize this and parallelize
+    // the work which would take longest.
+    for (int32_t outer = 0; outer < mDimProductOuter; outer++)
+    {
+        for (int32_t inner = 0; inner < mDimProductInner; inner++)
+        {
+            int32_t const axesOffset = outer * mDimProductInner * mAxisSize + inner;
+            float const* arr = &data[axesOffset];
+            int32_t const stride = mDimProductInner;
+            int32_t argmaxResult;
+            CUBLAS_CALL(cublasIsamax(mCublas, mAxisSize, arr, stride, &argmaxResult));
+
+            // cublasIsamax returns 1-indexed so convert to 0-indexed
+            argmaxResult--;
+
+            // cublasIsamax returns the index of the element with the highest absolute value.
+            // If this element is positive, then we know it is also the max.
+            // However, if it is negative, we need to
+            //      1) Copy the axis into our workspace
+            //      2) Subtract the minimum value we found from our array. This ensures that
+            //         none of the values are negative, and that the largest element remains
+            //         the largest element.
+            //      3) Use cublasIsamax to find the largest element again.
+            // NOTE: We are using cudaMemcpy instead of cudaMemcpyAsync because we need to know
+            //       maxAbsValue before proceeding. However, using synchronous rather than
+            //       asynchronous calls inside of enqueue() hurts performance.
+            //       This could be fixed by implementing the functionality of this plugin with a kernel
+            //       instead of relying only on cuBLAS.
+            float maxAbsValue;
+            CUDA_CALL(cudaMemcpy(&maxAbsValue, &arr[argmaxResult * stride], sizeof(float), cudaMemcpyDeviceToHost));
+            if (maxAbsValue < 0)
+            {
+                float negMinValue = -maxAbsValue;
+                CUBLAS_CALL(cublasScopy(mCublas, mAxisSize, arr, stride, axisFlat, 1));
+                CUBLAS_CALL(cublasSaxpy(mCublas, mAxisSize, &negMinValue, ones, 1, axisFlat, 1));
+                CUBLAS_CALL(cublasIsamax(mCublas, mAxisSize, axisFlat, 1, &argmaxResult));
+                argmaxResult--;
+            }
+
+            CUDA_CALL(cudaMemcpyAsync(
+                &result[axesOffset + argmaxResult * stride], &one, sizeof(float), cudaMemcpyHostToDevice, stream));
+        }
+    }
+    return cudaPeekAtLastError();
+}
+
+IPluginV3* HardmaxPlugin::attachToContext(IPluginResourceContext* context) noexcept
+{
+    auto* cloned = static_cast<HardmaxPlugin*>(clone());
+    if (cloned == nullptr)
+    {
+        return nullptr;
+    }
+    cublasStatus_t ret = cublasCreate(&cloned->mCublas);
+    ASSERT(ret == CUBLAS_STATUS_SUCCESS && cloned->mCublas != nullptr && "Failed to create cublasHandle_t.");
+    return cloned;
+}
+
+PluginFieldCollection const* HardmaxPlugin::getFieldsToSerialize() noexcept
+{
+    mDataToSerialize.clear();
+    mDataToSerialize.emplace_back("axis", &mAxis, PluginFieldType::kINT32, 1);
+    mFCToSerialize.nbFields = mDataToSerialize.size();
+    mFCToSerialize.fields = mDataToSerialize.data();
+    return &mFCToSerialize;
+}
+
+void HardmaxPlugin::setPluginNamespace(char const* libNamespace) noexcept
+{
+    ASSERT(libNamespace != nullptr);
+    mNamespace = libNamespace;
+}
+
+// HardmaxPluginCreator methods
+
+HardmaxPluginCreator::HardmaxPluginCreator()
+{
+    mPluginAttributes.clear();
+
+    // Consistent with the ONNX model attr fields
+    static auto const axisField = PluginField("axis", nullptr, PluginFieldType::kINT32, 1);
+    mPluginAttributes.emplace_back(axisField);
+
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+char const* HardmaxPluginCreator::getPluginName() const noexcept
+{
+    return kHARDMAX_NAME;
+}
+
+char const* HardmaxPluginCreator::getPluginVersion() const noexcept
+{
+    return kHARDMAX_VERSION;
+}
+
+PluginFieldCollection const* HardmaxPluginCreator::getFieldNames() noexcept
+{
+    return &mFC;
+}
+
+char const* HardmaxPluginCreator::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+void HardmaxPluginCreator::setPluginNamespace(char const* libNamespace) noexcept
+{
+    ASSERT(libNamespace != nullptr);
+    mNamespace = libNamespace;
+}
+
+IPluginV3* HardmaxPluginCreator::createPlugin(
+    char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
+{
+    using namespace std::string_view_literals;
+    // Set default value
+    int32_t axis = -1;
+
+    for (int32_t i = 0; i < fc->nbFields; i++)
+    {
+        if (fc->fields[i].name == "axis"sv)
+        {
+            ASSERT(fc->fields[i].type == PluginFieldType::kINT32);
+            axis = *static_cast<int32_t const*>(fc->fields[i].data);
+        }
+    }
+
+    auto plugin = std::make_unique<HardmaxPlugin>(axis);
+    plugin->setPluginNamespace(mNamespace.c_str());
+
+    return plugin.release();
+}
@@ -0,0 +1,131 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_HARDMAX_PLUGIN_H
+#define TRT_HARDMAX_PLUGIN_H
+
+#include "NvInferPlugin.h"
+#include <cublas_v2.h>
+#include <string>
+#include <vector>
+
+// This sample demonstrates how to implement a TensorRT plugin using the IPluginV3 interface.
+
+class HardmaxPlugin final : public nvinfer1::IPluginV3,
+                            public nvinfer1::IPluginV3OneCore,
+                            public nvinfer1::IPluginV3OneBuild,
+                            public nvinfer1::IPluginV3OneRuntime
+{
+public:
+    HardmaxPlugin() = delete;
+    HardmaxPlugin(int32_t axis);
+    HardmaxPlugin(HardmaxPlugin const& other);
+    ~HardmaxPlugin() override;
+
+    // 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* getPluginVersion() const noexcept override;
+    char const* getPluginNamespace() const noexcept override;
+
+    // IPluginV3OneBuild methods
+    int32_t getNbOutputs() const noexcept override;
+
+    int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
+        nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
+
+    bool supportsFormatCombination(int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs,
+        int32_t nbOutputs) 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;
+
+    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;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept;
+
+private:
+    std::string mNamespace;
+
+    // Number of elements in the axis along which hardmax is performed.
+    int32_t mAxisSize{0};
+
+    // Product of dimensions before and after mAxis.
+    // For example, if the input dimensions are [3, 4, 5, 6, 7] and mAxis = 2,
+    // then mDimProductOuter = 12 and mDimProductInner = 42.
+    int32_t mDimProductOuter{1};
+    int32_t mDimProductInner{1};
+
+    cublasHandle_t mCublas{nullptr};
+
+    // Attributes
+    // Axis along which to perform hardmax.
+    // Can be negative initially, but once configurePlugin() is called it will
+    // be converted to a positive axis.
+    int32_t mAxis{-1};
+
+    // Serialization helpers
+    std::vector<nvinfer1::PluginField> mDataToSerialize;
+    nvinfer1::PluginFieldCollection mFCToSerialize;
+};
+
+class HardmaxPluginCreator : public nvinfer1::IPluginCreatorV3One
+{
+public:
+    HardmaxPluginCreator();
+
+    ~HardmaxPluginCreator() override = default;
+
+    char const* getPluginName() const noexcept override;
+
+    char const* getPluginVersion() const noexcept override;
+
+    nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
+
+    nvinfer1::IPluginV3* createPlugin(
+        char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept;
+
+    char const* getPluginNamespace() const noexcept override;
+
+private:
+    nvinfer1::PluginFieldCollection mFC;
+    std::vector<nvinfer1::PluginField> mPluginAttributes;
+    std::string mNamespace;
+};
+
+#endif // TRT_HARDMAX_PLUGIN_H
@@ -0,0 +1,11 @@
+nltk==3.9.1
+onnx==1.18.0
+--extra-index-url https://pypi.ngc.nvidia.com
+onnx-graphsurgeon>=0.3.20
+wget>=3.2
+cuda-python==12.9.0
+pywin32; platform_system == "Windows"
+pyyaml==6.0.3
+requests==2.32.4
+tqdm==4.66.4
+numpy==1.26.4
@@ -0,0 +1,193 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import sys
+
+import tensorrt as trt
+
+from model import TRT_MODEL_PATH
+from load_plugin_lib import load_plugin_lib
+
+# ../common.py
+parent_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)
+sys.path.insert(1, parent_dir)
+import common
+
+# Reuse some BiDAF-specific methods
+# ../engine_refit_onnx_bidaf/data_processing.py
+sys.path.insert(1, os.path.join(parent_dir, "engine_refit_onnx_bidaf"))
+from engine_refit_onnx_bidaf.data_processing import preprocess, get_inputs
+
+# Maxmimum number of words in context or query text.
+# Used in optimization profile when building engine.
+# Adjustable.
+MAX_TEXT_LENGTH = 64
+
+WORKING_DIR = os.environ.get("TRT_WORKING_DIR") or os.path.dirname(
+    os.path.realpath(__file__)
+)
+
+# Path to which trained model will be saved (check README.md)
+ENGINE_FILE_PATH = os.path.join(WORKING_DIR, "bidaf.trt")
+
+# Define global logger object (it should be a singleton,
+# available for TensorRT from anywhere in code).
+# You can set the logger severity higher to suppress messages
+# (or lower to display more messages)
+TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+
+
+# Builds TensorRT Engine
+def build_engine(model_path):
+
+    builder = trt.Builder(TRT_LOGGER)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    config = builder.create_builder_config()
+    parser = trt.OnnxParser(network, TRT_LOGGER)
+    runtime = trt.Runtime(TRT_LOGGER)
+
+    # Parse model file
+    print("Loading ONNX file from path {}...".format(model_path))
+    with open(model_path, "rb") as model:
+        print("Beginning ONNX file parsing")
+        if not parser.parse(model.read()):
+            print("ERROR: Failed to parse the ONNX file.")
+            for error in range(parser.num_errors):
+                print(parser.get_error(error))
+            return None
+    print("Completed parsing of ONNX file")
+
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, common.GiB(1))
+
+    # The input text length is variable, so we need to specify an optimization profile.
+    profile = builder.create_optimization_profile()
+    for i in range(network.num_inputs):
+        input = network.get_input(i)
+        assert input.shape[0] == -1
+        min_shape = [1] + list(input.shape[1:])
+        opt_shape = [8] + list(input.shape[1:])
+        max_shape = [MAX_TEXT_LENGTH] + list(input.shape[1:])
+        profile.set_shape(input.name, min_shape, opt_shape, max_shape)
+    config.add_optimization_profile(profile)
+
+    print("Building TensorRT engine. This may take a few minutes.")
+    plan = builder.build_serialized_network(network, config)
+    engine = runtime.deserialize_cuda_engine(plan)
+    with open(ENGINE_FILE_PATH, "wb") as f:
+        f.write(plan)
+    return engine
+
+
+def load_test_case(inputs, context_text, query_text, trt_context):
+    # Part 1: Specify Input shapes
+    cw, cc = preprocess(context_text)
+    qw, qc = preprocess(query_text)
+    for arr in (cw, cc, qw, qc):
+        assert arr.shape[0] <= MAX_TEXT_LENGTH, (
+            "Input context or query is too long! "
+            + "Either decrease the input length or increase MAX_TEXT_LENGTH"
+        )
+    trt_context.set_input_shape("CategoryMapper_4", cw.shape)
+    trt_context.set_input_shape("CategoryMapper_5", cc.shape)
+    trt_context.set_input_shape("CategoryMapper_6", qw.shape)
+    trt_context.set_input_shape("CategoryMapper_7", qc.shape)
+
+    # Part 2: load input data
+    cw_flat, cc_flat, qw_flat, qc_flat = get_inputs(context_text, query_text)
+    for i, arr in enumerate([cw_flat, cc_flat, qw_flat, qc_flat]):
+        inputs[i].host = arr
+
+
+def main():
+    # Load the shared object file containing the Hardmax plugin implementation.
+    # By doing this, you will also register the Hardmax plugin with the TensorRT
+    # PluginRegistry through use of the macro REGISTER_TENSORRT_PLUGIN present
+    # in the plugin implementation. Refer to plugin/customHardmaxPlugin.cpp for more details.
+    load_plugin_lib()
+
+    # Load pretrained model
+    if not os.path.isfile(TRT_MODEL_PATH):
+        raise IOError(
+            "\n{}\n{}\n{}\n".format(
+                "Failed to load model file ({}).".format(TRT_MODEL_PATH),
+                "Please use 'python3 model.py' to generate the ONNX model.",
+                "For more information, see README.md",
+            )
+        )
+
+    if os.path.exists(ENGINE_FILE_PATH):
+        print(f"Loading saved TRT engine from {ENGINE_FILE_PATH}")
+        with open(ENGINE_FILE_PATH, "rb") as f:
+            runtime = trt.Runtime(TRT_LOGGER)
+            runtime.max_threads = 10
+            engine = runtime.deserialize_cuda_engine(f.read())
+    else:
+        print("Engine plan not saved. Building new engine...")
+        engine = build_engine(TRT_MODEL_PATH)
+
+    inputs, outputs, bindings = common.allocate_buffers(engine, profile_idx=0)
+
+    testcases = [
+        (
+            "Garry the lion is 5 years old. He lives in the savanna.",
+            "Where does the lion live?",
+        ),
+        ("A quick brown fox jumps over the lazy dog.", "What color is the fox?"),
+    ]
+
+    print("\n=== Testing ===")
+
+    interactive = "--interactive" in sys.argv
+    if interactive:
+        context_text = input("Enter context: ")
+        query_text = input("Enter query: ")
+        testcases = [(context_text, query_text)]
+
+    trt_context = engine.create_execution_context()
+    
+    # Use context manager for proper stream lifecycle management
+    with common.CudaStreamContext() as stream:
+        for context_text, query_text in testcases:
+
+            context_words, _ = preprocess(context_text)
+
+            load_test_case(inputs, context_text, query_text, trt_context)
+            if not interactive:
+                print(f"Input context: {context_text}")
+                print(f"Input query: {query_text}")
+            trt_outputs = common.do_inference(
+                trt_context,
+                engine=engine,
+                bindings=bindings,
+                inputs=inputs,
+                outputs=outputs,
+                stream=stream,
+            )
+            start = trt_outputs[1].item()
+            end = trt_outputs[0].item()
+            answer = context_words[start : end + 1].flatten()
+            print(f"Model prediction: ", " ".join(answer))
+            print()
+        
+        # Note: free_buffers no longer needs stream parameter
+        common.free_buffers(inputs, outputs)
+    print("Passed")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,102 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import numpy as np
+import os
+import sys
+import tensorrt as trt
+
+# ../common.py
+parent_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)
+sys.path.insert(1, parent_dir)
+import common
+
+from load_plugin_lib import load_plugin_lib
+
+TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
+
+
+def hardmax_reference_impl(arr, axis):
+    one_hot = np.zeros(arr.shape, dtype=arr.dtype)
+    argmax = np.expand_dims(np.argmax(arr, axis), axis)
+    np.put_along_axis(one_hot, argmax, 1, axis=axis)
+    return one_hot
+
+
+def make_trt_network_and_engine(input_shape, axis):
+    registry = trt.get_plugin_registry()
+    plugin_creator = registry.get_creator("CustomHardmax", "1", "")
+    axis_buffer = np.array([axis])
+    axis_attr = trt.PluginField("axis", axis_buffer, type=trt.PluginFieldType.INT32)
+    field_collection = trt.PluginFieldCollection([axis_attr])
+    plugin = plugin_creator.create_plugin(
+        name="CustomHardmax", field_collection=field_collection, phase=trt.TensorRTPhase.BUILD
+    )
+
+    builder = trt.Builder(TRT_LOGGER)
+    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    config = builder.create_builder_config()
+    runtime = trt.Runtime(TRT_LOGGER)
+
+    input_layer = network.add_input(
+        name="input_layer", dtype=trt.float32, shape=input_shape
+    )
+    hardmax = network.add_plugin_v3(inputs=[input_layer], shape_inputs=[], plugin=plugin)
+    network.mark_output(hardmax.get_output(0))
+
+    plan = builder.build_serialized_network(network, config)
+    engine = runtime.deserialize_cuda_engine(plan)
+
+    return engine
+
+
+def custom_plugin_impl(input_arr, engine):
+    inputs, outputs, bindings = common.allocate_buffers(engine)
+    context = engine.create_execution_context()
+    inputs[0].host = input_arr.astype(trt.nptype(trt.float32))
+    
+    # Use context manager for proper stream lifecycle management
+    with common.CudaStreamContext() as stream:
+        trt_outputs = common.do_inference(
+            context,
+            engine=engine,
+            bindings=bindings,
+            inputs=inputs,
+            outputs=outputs,
+            stream=stream,
+        )
+        output = trt_outputs[0].copy()
+    common.free_buffers(inputs, outputs)
+    return output
+
+
+def main():
+    load_plugin_lib()
+    for num_dims in range(1, 8):
+        for axis in range(-num_dims, num_dims):
+            shape = np.random.randint(1, 4, size=num_dims)
+            arr = np.random.rand(*shape)
+            arr = (arr - 0.5) * 200
+            engine = make_trt_network_and_engine(shape, axis)
+            res1 = hardmax_reference_impl(arr, axis)
+            res2 = custom_plugin_impl(arr, engine).reshape(res1.shape)
+            assert np.all(res1 == res2), f"Test failed for shape={shape}, axis={axis}"
+    print("Passed")
+
+
+if __name__ == "__main__":
+    main()