chore: import upstream snapshot with attribution

plugin/topkLastDimPlugin/CMakeLists.txt

@@ -0,0 +1,23 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 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.
+#
+
+add_plugin_source(
+    topkLastDimPlugin.cpp
+    topkLastDimPlugin.h
+    transpose.cu
+    transpose.h
+)

plugin/topkLastDimPlugin/README.md

@@ -0,0 +1,62 @@
+# TopkLastDim
+
+**Table Of Contents**
+- [Description](#description)
+    * [Structure](#structure)
+- [Parameters](#parameters)
+- [Additional resources](#additional-resources)
+- [License](#license)
+- [Changelog](#changelog)
+- [Known issues](#known-issues)
+
+## Description
+
+The `TopkLastDim` plugin computes the top-`k` largest or smallest elements of an input tensor along a specified axis, following the [ONNX specification for TopK](https://github.com/onnx/onnx/blob/main/docs/Operators.md#TopK) for value selection, ordering, and the meaning of the `axis` / `k` / `largest` parameters. The plugin emits `int32` indices rather than the `int64` indices required by the ONNX spec — see [Known issues](#known-issues).
+
+The plugin is built around the AIR (Adaptive Iterative Radix) sort kernel ported from [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/topkLastDim.cu), which operates on the last dimension of a 2D `[numRows, rowLength]` view. The plugin handles arbitrary input rank and arbitrary axis by:
+
+- **Fast path** (`axis == last dim`): the kernel is invoked directly on the input with no copies.
+- **General path** (any other axis): the input is transposed so the target axis is moved to the last position, the kernel runs on the resulting 2D view, and the output values and indices are transposed back so the top-`k` dimension sits at the original axis.
+
+### Structure
+
+The `TopkLastDim` plugin consumes the following input:
+
+1. `input` - T: A tensor of arbitrary rank. The plugin computes top-`k` along the axis specified by the `axis` attribute. T can be `float32`, `float16`, `int32`, or `bfloat16`.
+
+The `TopkLastDim` plugin produces the following outputs:
+
+1. `values` - T: A tensor with the same shape as `input` except that the size along `axis` is replaced by `k`. Contains the top-`k` values along `axis`, in descending order when `is_largest == 1` and ascending order when `is_largest == 0`.
+2. `indices` - `int32`: A tensor with the same shape as `values`. Contains the indices, along `axis` of `input`, of the corresponding entries in `values`. Note: ONNX TopK specifies `int64` indices; this plugin emits `int32`.
+
+This plugin has the plugin creator class `TopkLastDimPluginCreator` and the plugin class `TopkLastDimPlugin` which extends `IPluginV3`.
+
+## Parameters
+
+The `TopkLastDim` plugin has the following parameters:
+
+| Type    | Parameter     | Description
+|---------|---------------|--------------------------------------------------------
+| `int32` | `type_id`     | Data type of the input tensor (and of the `values` output). Allowed values follow `nvinfer1::DataType`: `0` (kFLOAT), `1` (kHALF), `3` (kINT32), `7` (kBF16). Required.
+| `int32` | `k`           | Number of top elements to return along `axis`. Must be a positive integer not greater than the size of `input` along `axis`. Required.
+| `int32` | `is_largest`  | If `1`, return the `k` largest elements (sorted descending). If `0`, return the `k` smallest elements (sorted ascending). Required.
+| `int32` | `axis`        | Axis along which to compute top-`k`. Negative values count from the back (e.g., `-1` means the last dimension). Optional; default is `-1`.
+
+## Additional resources
+
+The following resources provide a deeper understanding of the `TopkLastDim` plugin:
+
+- [ONNX specification for TopK](https://github.com/onnx/onnx/blob/main/docs/Operators.md#TopK)
+- [AIR TopK kernel in TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/topkLastDim.cu)
+
+## 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).
+
+## Changelog
+
+May 2026: This is the first release of this `README.md` file.
+
+## Known issues
+
+- ONNX TopK specifies indices of type `int64`; this plugin emits `int32` indices to match the upstream TRT-LLM kernel's public entry point (which fixes `IdxT = int32_t`). Inputs along `axis` longer than `2^31 - 1` are therefore not addressable.

plugin/topkLastDimPlugin/TopkLastDim_PluginConfig.yaml

@@ -0,0 +1,110 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 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.
+#
+---
+name: TopkLastDim
+versions:
+  "1":
+    interface: "IPluginV3"
+    inputs:
+      - input
+    outputs:
+      - values
+      - indices
+    attributes:
+      - type_id
+      - k
+      - is_largest
+      - axis
+    attribute_types:
+      type_id: int32
+      k: int32
+      is_largest: int32
+      axis: int32
+    attribute_length:
+      type_id: 1
+      k: 1
+      is_largest: 1
+      axis: 1
+    attribute_options:
+      type_id:
+        - 0   # kFLOAT
+        - 1   # kHALF
+        - 3   # kINT32
+        - 7   # kBF16
+      k:
+        min: "1"
+        max: "=pinf"
+      is_largest:
+        - 0
+        - 1
+      axis:
+        min: "=ninf"
+        max: "=pinf"
+    attributes_required:
+      - type_id
+      - k
+      - is_largest
+    configs:
+      config_float:
+        input_types:
+          input: float32
+        attribute_options:
+          type_id:
+            value: 0
+          k:
+            value: 5
+          is_largest:
+            value: 1
+          axis:
+            value: -1
+      config_half:
+        input_types:
+          input: float16
+        attribute_options:
+          type_id:
+            value: 1
+          k:
+            value: 5
+          is_largest:
+            value: 1
+          axis:
+            value: -1
+      config_int32:
+        input_types:
+          input: int32
+        attribute_options:
+          type_id:
+            value: 3
+          k:
+            value: 5
+          is_largest:
+            value: 1
+          axis:
+            value: -1
+      config_bf16:
+        input_types:
+          input: bfloat16
+        attribute_options:
+          type_id:
+            value: 7
+          k:
+            value: 5
+          is_largest:
+            value: 1
+          axis:
+            value: -1
+...

plugin/topkLastDimPlugin/topkLastDimPlugin.cpp

@@ -0,0 +1,526 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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 "topkLastDimPlugin.h"
+#include "common/checkMacrosPlugin.h"
+#include "common/plugin.h"
+#include "transpose.h"
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <limits>
+#include <mutex>
+#include <string_view>
+
+namespace nvinfer1::plugin
+{
+
+// Kernel API implemented in topkLastDim.cu.
+// The __restrict__ qualifiers must match the definition in topkLastDim.cu exactly,
+// because MSVC includes __restrict in the mangled symbol name.
+template <typename T>
+size_t invokeComputeTopkLastDimWorkspaceSize(int32_t batchSize, int32_t inputLength, int32_t k, bool is_largest);
+template <typename T>
+void invokeTopkLastDim(int32_t batchSize, int32_t inputLength, int32_t k, bool is_largest,
+    void const* __restrict__ input, void* __restrict__ out_val, void* __restrict__ out_idx, void* workspace,
+    cudaStream_t stream);
+
+namespace
+{
+char const* gKTopkLastDimPluginVersion{"1"};
+char const* gKTopkLastDimPluginName{"TopkLastDim"};
+} // namespace
+
+// ========================== Plugin ==========================
+
+TopkLastDimPlugin::TopkLastDimPlugin(int32_t typeId, int32_t k, int32_t isLargest, int32_t axis)
+    : mTypeId(typeId)
+    , mK(k)
+    , mIsLargest(isLargest)
+    , mAxis(axis)
+{
+    auto const type = static_cast<DataType>(mTypeId);
+    PLUGIN_VALIDATE(
+        type == DataType::kBF16 || type == DataType::kFLOAT || type == DataType::kHALF || type == DataType::kINT32);
+    PLUGIN_VALIDATE(mK > 0);
+    PLUGIN_VALIDATE(mIsLargest == 0 || mIsLargest == 1);
+}
+
+int32_t TopkLastDimPlugin::resolveAxis(int32_t nbDims) const
+{
+    int32_t axis = mAxis;
+    if (axis < 0)
+    {
+        axis += nbDims;
+    }
+    PLUGIN_ASSERT(axis >= 0 && axis < nbDims);
+    return axis;
+}
+
+int32_t TopkLastDimPlugin::elementSize() const
+{
+    auto const type = static_cast<DataType>(mTypeId);
+    switch (type)
+    {
+    case DataType::kFLOAT:
+    case DataType::kINT32: return 4;
+    case DataType::kHALF:
+    case DataType::kBF16: return 2;
+    default: PLUGIN_ASSERT(false && "Unsupported data type"); return 0;
+    }
+}
+
+size_t TopkLastDimPlugin::topkKernelWorkspaceSize(int32_t numRows, int32_t rowLength) const
+{
+    bool const isLargest = mIsLargest != 0;
+    auto const type = static_cast<DataType>(mTypeId);
+    if (type == DataType::kINT32)
+    {
+        return invokeComputeTopkLastDimWorkspaceSize<int32_t>(numRows, rowLength, mK, isLargest);
+    }
+    if (type == DataType::kHALF)
+    {
+        return invokeComputeTopkLastDimWorkspaceSize<half>(numRows, rowLength, mK, isLargest);
+    }
+    if (type == DataType::kFLOAT)
+    {
+        return invokeComputeTopkLastDimWorkspaceSize<float>(numRows, rowLength, mK, isLargest);
+    }
+    if (type == DataType::kBF16)
+    {
+        return invokeComputeTopkLastDimWorkspaceSize<__nv_bfloat16>(numRows, rowLength, mK, isLargest);
+    }
+    PLUGIN_ASSERT(false && "Unsupported data type");
+    return 0;
+}
+
+size_t TopkLastDimPlugin::transposeWorkspaceSize(Dims const& dims) const
+{
+    int32_t const nbDims = dims.nbDims;
+    int32_t const axis = resolveAxis(nbDims);
+
+    // No transpose needed when axis is already the last dimension.
+    if (axis == nbDims - 1)
+    {
+        return 0;
+    }
+
+    int64_t totalElems = 1;
+    for (int32_t i = 0; i < nbDims; ++i)
+    {
+        totalElems *= dims.d[i];
+    }
+
+    int32_t const elemSz = elementSize();
+    // Transposed input buffer (values).
+    size_t bytes = totalElems * elemSz;
+    // Transposed output values buffer.
+    int64_t const axisLen = dims.d[axis];
+    if (axisLen == 0)
+    {
+        return 0;
+    }
+    int64_t const outputElems = (totalElems / axisLen) * mK;
+    bytes += outputElems * elemSz;
+    // Alignment padding so transposedIndices (int32_t*) is 4-byte aligned.
+    bytes = (bytes + alignof(int32_t) - 1) & ~(alignof(int32_t) - 1);
+    // Transposed output indices buffer.
+    bytes += outputElems * sizeof(int32_t);
+    return bytes;
+}
+
+IPluginCapability* TopkLastDimPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept
+{
+    try
+    {
+        if (type == PluginCapabilityType::kBUILD)
+        {
+            return static_cast<IPluginV3OneBuild*>(this);
+        }
+        if (type == PluginCapabilityType::kRUNTIME)
+        {
+            return static_cast<IPluginV3OneRuntime*>(this);
+        }
+        PLUGIN_ASSERT(type == PluginCapabilityType::kCORE);
+        return static_cast<IPluginV3OneCore*>(this);
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+IPluginV3* TopkLastDimPlugin::clone() noexcept
+{
+    try
+    {
+        auto plugin = std::make_unique<TopkLastDimPlugin>(*this);
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+char const* TopkLastDimPlugin::getPluginName() const noexcept
+{
+    return gKTopkLastDimPluginName;
+}
+
+char const* TopkLastDimPlugin::getPluginVersion() const noexcept
+{
+    return gKTopkLastDimPluginVersion;
+}
+
+char const* TopkLastDimPlugin::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+int32_t TopkLastDimPlugin::getNbOutputs() const noexcept
+{
+    return 2;
+}
+
+int32_t TopkLastDimPlugin::configurePlugin(
+    DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    return 0;
+}
+
+bool TopkLastDimPlugin::supportsFormatCombination(
+    int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
+{
+    PLUGIN_ASSERT(inOut != nullptr);
+    PLUGIN_ASSERT(pos >= 0 && pos <= 2);
+    PLUGIN_ASSERT(nbInputs == 1);
+    PLUGIN_ASSERT(nbOutputs == 2);
+
+    PluginTensorDesc const& desc = inOut[pos].desc;
+    if (desc.format != TensorFormat::kLINEAR)
+    {
+        return false;
+    }
+    // Input (pos 0) and values output (pos 1) must match the configured type.
+    if (pos < 2)
+    {
+        return desc.type == static_cast<DataType>(mTypeId);
+    }
+    // Indices output (pos 2) is always INT32.
+    return desc.type == DataType::kINT32;
+}
+
+int32_t TopkLastDimPlugin::getOutputDataTypes(
+    DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
+{
+    PLUGIN_ASSERT(nbInputs == 1);
+    PLUGIN_ASSERT(nbOutputs == 2);
+    outputTypes[0] = inputTypes[0];    // values: same type as input
+    outputTypes[1] = DataType::kINT32; // indices
+    return 0;
+}
+
+int32_t TopkLastDimPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
+    int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept
+{
+    PLUGIN_ASSERT(nbInputs == 1);
+    PLUGIN_ASSERT(nbOutputs == 2);
+
+    int32_t const nbDims = inputs[0].nbDims;
+    int32_t const axis = resolveAxis(nbDims);
+
+    auto const* kExpr = exprBuilder.constant(mK);
+    PLUGIN_ASSERT(kExpr != nullptr);
+
+    // Output shape is same as input but with dim[axis] replaced by k.
+    for (int32_t o = 0; o < 2; ++o)
+    {
+        outputs[o] = inputs[0];
+        outputs[o].d[axis] = kExpr;
+    }
+    return 0;
+}
+
+size_t TopkLastDimPlugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+    DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
+{
+    Dims const& maxDims = inputs[0].max;
+    int32_t const nbDims = maxDims.nbDims;
+    int32_t const axis = resolveAxis(nbDims);
+
+    // Compute the 2D shape the kernel will see.
+    int64_t numRows = 1;
+    for (int32_t i = 0; i < nbDims; ++i)
+    {
+        if (i != axis)
+        {
+            numRows *= maxDims.d[i];
+        }
+    }
+    int32_t const rowLength = maxDims.d[axis];
+
+    PLUGIN_ASSERT(numRows <= std::numeric_limits<int32_t>::max());
+    size_t bytes = topkKernelWorkspaceSize(static_cast<int32_t>(numRows), rowLength);
+    bytes += transposeWorkspaceSize(maxDims);
+    return bytes;
+}
+
+template <typename T>
+int32_t TopkLastDimPlugin::enqueueImpl(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
+    void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream)
+{
+    Dims const& dims = inputDesc[0].dims;
+    int32_t const nbDims = dims.nbDims;
+    int32_t const axis = resolveAxis(nbDims);
+
+    // Compute outer, axisLen, inner for the 3D view [outer, axisLen, inner].
+    int64_t outer = 1;
+    for (int32_t i = 0; i < axis; ++i)
+    {
+        outer *= dims.d[i];
+    }
+    int32_t const axisLen = dims.d[axis];
+    if (axisLen == 0)
+    {
+        return 0; // Empty tensor along axis dimension
+    }
+    int64_t inner = 1;
+    for (int32_t i = axis + 1; i < nbDims; ++i)
+    {
+        inner *= dims.d[i];
+    }
+
+    PLUGIN_ASSERT(outer * inner <= std::numeric_limits<int32_t>::max());
+    int32_t const numRows = static_cast<int32_t>(outer * inner);
+    if (numRows == 0)
+    {
+        return 0;
+    }
+
+    bool const isLargest = mIsLargest != 0;
+
+    // Fast path: axis is already the last dimension — call the kernel directly.
+    if (axis == nbDims - 1)
+    {
+        invokeTopkLastDim<T>(numRows, axisLen, mK, isLargest, inputs[0], outputs[0], outputs[1], workspace, stream);
+        PLUGIN_CUASSERT(cudaGetLastError());
+        return 0;
+    }
+
+    // Multi-dimensional tensor path: transpose -> topk -> transpose back.
+    int32_t const elemSz = sizeof(T);
+    int64_t const totalInputElems = outer * axisLen * inner;
+    int64_t const totalOutputElems = outer * inner * mK;
+
+    // Partition workspace: [transposedInput | transposedValues | transposedIndices | topkWorkspace]
+    char* wsPtr = static_cast<char*>(workspace);
+    T* transposedInput = reinterpret_cast<T*>(wsPtr);
+    wsPtr += totalInputElems * elemSz;
+    T* transposedValues = reinterpret_cast<T*>(wsPtr);
+    wsPtr += totalOutputElems * elemSz;
+    // Align to 4 bytes for int32_t (needed when T is a 2-byte type and element count is odd).
+    auto aligned = (reinterpret_cast<uintptr_t>(wsPtr) + alignof(int32_t) - 1) & ~(alignof(int32_t) - 1);
+    int32_t* transposedIndices = reinterpret_cast<int32_t*>(aligned);
+    wsPtr = reinterpret_cast<char*>(aligned);
+    wsPtr += totalOutputElems * sizeof(int32_t);
+    void* topkWorkspace = wsPtr;
+
+    // Step 1: Transpose input [outer, axisLen, inner] -> [outer, inner, axisLen]
+    launchBatchedTranspose2D<T>(static_cast<T const*>(inputs[0]), transposedInput, static_cast<int32_t>(outer), axisLen,
+        static_cast<int32_t>(inner), stream);
+
+    // Step 2: Run TopK on the 2D view [outer*inner, axisLen]
+    invokeTopkLastDim<T>(
+        numRows, axisLen, mK, isLargest, transposedInput, transposedValues, transposedIndices, topkWorkspace, stream);
+
+    // Step 3: Transpose outputs [outer, inner, K] -> [outer, K, inner]
+    launchBatchedTranspose2D<T>(transposedValues, static_cast<T*>(outputs[0]), static_cast<int32_t>(outer),
+        static_cast<int32_t>(inner), mK, stream);
+    launchBatchedTranspose2D<int32_t>(transposedIndices, static_cast<int32_t*>(outputs[1]), static_cast<int32_t>(outer),
+        static_cast<int32_t>(inner), mK, stream);
+
+    PLUGIN_CUASSERT(cudaGetLastError());
+    return 0;
+}
+
+int32_t TopkLastDimPlugin::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
+    void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
+{
+    auto const type = static_cast<DataType>(mTypeId);
+    if (type == DataType::kINT32)
+    {
+        return enqueueImpl<int32_t>(inputDesc, outputDesc, inputs, outputs, workspace, stream);
+    }
+    if (type == DataType::kHALF)
+    {
+        return enqueueImpl<half>(inputDesc, outputDesc, inputs, outputs, workspace, stream);
+    }
+    if (type == DataType::kFLOAT)
+    {
+        return enqueueImpl<float>(inputDesc, outputDesc, inputs, outputs, workspace, stream);
+    }
+    if (type == DataType::kBF16)
+    {
+        return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream);
+    }
+    PLUGIN_ASSERT(false && "Unsupported data type");
+    return 0;
+}
+
+int32_t TopkLastDimPlugin::onShapeChange(
+    PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    return 0;
+}
+
+IPluginV3* TopkLastDimPlugin::attachToContext(IPluginResourceContext* context) noexcept
+{
+    return clone();
+}
+
+PluginFieldCollection const* TopkLastDimPlugin::getFieldsToSerialize() noexcept
+{
+    mDataToSerialize.clear();
+    mDataToSerialize.emplace_back("type_id", &mTypeId, PluginFieldType::kINT32, 1);
+    mDataToSerialize.emplace_back("k", &mK, PluginFieldType::kINT32, 1);
+    mDataToSerialize.emplace_back("is_largest", &mIsLargest, PluginFieldType::kINT32, 1);
+    mDataToSerialize.emplace_back("axis", &mAxis, PluginFieldType::kINT32, 1);
+    mFCToSerialize.nbFields = mDataToSerialize.size();
+    mFCToSerialize.fields = mDataToSerialize.data();
+    return &mFCToSerialize;
+}
+
+void TopkLastDimPlugin::setPluginNamespace(char const* pluginNamespace) noexcept
+{
+    try
+    {
+        PLUGIN_ASSERT(pluginNamespace != nullptr);
+        mNamespace = pluginNamespace;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+// ========================== Creator ==========================
+
+TopkLastDimPluginCreator::TopkLastDimPluginCreator()
+{
+    static std::mutex sMutex;
+    std::lock_guard<std::mutex> guard(sMutex);
+    mPluginAttributes.clear();
+    mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("k", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("is_largest", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("axis", nullptr, PluginFieldType::kINT32, 1));
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+char const* TopkLastDimPluginCreator::getPluginName() const noexcept
+{
+    return gKTopkLastDimPluginName;
+}
+
+char const* TopkLastDimPluginCreator::getPluginVersion() const noexcept
+{
+    return gKTopkLastDimPluginVersion;
+}
+
+PluginFieldCollection const* TopkLastDimPluginCreator::getFieldNames() noexcept
+{
+    return &mFC;
+}
+
+IPluginV3* TopkLastDimPluginCreator::createPlugin(
+    char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(fc != nullptr);
+        PluginField const* fields = fc->fields;
+
+        int32_t typeId{};
+        int32_t k{};
+        int32_t isLargest{};
+        int32_t axis{-1}; // default: last dimension
+        bool hasTypeId = false;
+        bool hasK = false;
+        bool hasIsLargest = false;
+
+        using namespace std::string_view_literals;
+        for (int32_t i = 0; i < fc->nbFields; ++i)
+        {
+            std::string_view const attrName = fields[i].name;
+            if (attrName == "type_id"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                typeId = *static_cast<int32_t const*>(fields[i].data);
+                hasTypeId = true;
+            }
+            else if (attrName == "k"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                k = *static_cast<int32_t const*>(fields[i].data);
+                hasK = true;
+            }
+            else if (attrName == "is_largest"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                isLargest = *static_cast<int32_t const*>(fields[i].data);
+                hasIsLargest = true;
+            }
+            else if (attrName == "axis"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                axis = *static_cast<int32_t const*>(fields[i].data);
+            }
+        }
+        PLUGIN_VALIDATE(hasTypeId, "Missing required field 'type_id'");
+        PLUGIN_VALIDATE(hasK, "Missing required field 'k'");
+        PLUGIN_VALIDATE(hasIsLargest, "Missing required field 'is_largest'");
+        return new TopkLastDimPlugin(typeId, k, isLargest, axis);
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+char const* TopkLastDimPluginCreator::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+void TopkLastDimPluginCreator::setPluginNamespace(char const* pluginNamespace) noexcept
+{
+    try
+    {
+        PLUGIN_ASSERT(pluginNamespace != nullptr);
+        mNamespace = pluginNamespace;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+} // namespace nvinfer1::plugin

plugin/topkLastDimPlugin/topkLastDimPlugin.h

@@ -0,0 +1,123 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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_TOPK_LAST_DIM_PLUGIN_H
+#define TRT_TOPK_LAST_DIM_PLUGIN_H
+
+#include "NvInfer.h"
+#include "common/plugin.h"
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace nvinfer1::plugin
+{
+
+class TopkLastDimPlugin : public IPluginV3,
+                          public IPluginV3OneCore,
+                          public IPluginV3OneBuild,
+                          public IPluginV3OneRuntime
+{
+public:
+    //! \param axis Axis along which to compute top-k. -1 means last dimension.
+    TopkLastDimPlugin(int32_t typeId, int32_t k, int32_t isLargest, int32_t axis = -1);
+    TopkLastDimPlugin(TopkLastDimPlugin const&) = default;
+    ~TopkLastDimPlugin() override = default;
+
+    // IPluginV3
+    IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override;
+    IPluginV3* clone() noexcept override;
+
+    // IPluginV3OneCore
+    char const* getPluginName() const noexcept override;
+    char const* getPluginVersion() const noexcept override;
+    char const* getPluginNamespace() const noexcept override;
+
+    // IPluginV3OneBuild
+    int32_t getNbOutputs() const noexcept override;
+    int32_t configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out,
+        int32_t nbOutputs) noexcept override;
+    bool supportsFormatCombination(
+        int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;
+    int32_t getOutputDataTypes(
+        DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override;
+    int32_t getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
+        int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept override;
+    size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+        DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
+
+    // IPluginV3OneRuntime
+    int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs,
+        void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
+    int32_t onShapeChange(
+        PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
+    IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override;
+    PluginFieldCollection const* getFieldsToSerialize() noexcept override;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept;
+
+private:
+    //! Resolve mAxis to a non-negative value given the input rank.
+    int32_t resolveAxis(int32_t nbDims) const;
+
+    //! Compute the element size for the configured data type.
+    int32_t elementSize() const;
+
+    //! Compute workspace needed for transpose buffers (0 when axis == last dim).
+    size_t transposeWorkspaceSize(Dims const& inputDims) const;
+
+    //! Compute workspace needed by the AIR TopK kernel itself.
+    size_t topkKernelWorkspaceSize(int32_t numRows, int32_t rowLength) const;
+
+    template <typename T>
+    int32_t enqueueImpl(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
+        void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream);
+
+    int32_t mTypeId; //!< DataType stored as int32_t for field-based serialization.
+    int32_t mK;
+    int32_t mIsLargest; //!< Boolean stored as int32_t for field-based serialization.
+    int32_t mAxis;      //!< Axis for top-k. -1 means last dimension.
+
+    std::string mNamespace;
+
+    std::vector<nvinfer1::PluginField> mDataToSerialize;
+    nvinfer1::PluginFieldCollection mFCToSerialize;
+};
+
+class TopkLastDimPluginCreator : public nvinfer1::IPluginCreatorV3One
+{
+public:
+    TopkLastDimPluginCreator();
+    ~TopkLastDimPluginCreator() override = default;
+
+    char const* getPluginName() const noexcept override;
+    char const* getPluginVersion() const noexcept override;
+    PluginFieldCollection const* getFieldNames() noexcept override;
+    IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
+    char const* getPluginNamespace() const noexcept override;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept;
+
+private:
+    PluginFieldCollection mFC;
+    std::vector<PluginField> mPluginAttributes;
+    std::string mNamespace;
+};
+
+} // namespace nvinfer1::plugin
+
+#endif // TRT_TOPK_LAST_DIM_PLUGIN_H

plugin/topkLastDimPlugin/transpose.cu

@@ -0,0 +1,71 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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 "transpose.h"
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime_api.h>
+
+namespace nvinfer1::plugin
+{
+
+/// Transpose the last two dimensions of a 3D tensor [outer, rows, cols] -> [outer, cols, rows].
+/// Each thread handles one element.
+template <typename T>
+__global__ void batchedTranspose2DKernel(
+    T const* __restrict__ src, T* __restrict__ dst, int32_t outer, int32_t rows, int32_t cols)
+{
+    int64_t const idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    int64_t const total = static_cast<int64_t>(outer) * rows * cols;
+    if (idx >= total)
+    {
+        return;
+    }
+
+    int64_t const matSize = static_cast<int64_t>(rows) * cols;
+    int32_t const o = static_cast<int32_t>(idx / matSize);
+    int32_t const rem = static_cast<int32_t>(idx % matSize);
+    int32_t const r = rem / cols;
+    int32_t const c = rem % cols;
+
+    // src[o][r][c] -> dst[o][c][r]
+    dst[static_cast<int64_t>(o) * matSize + static_cast<int64_t>(c) * rows + r] = src[idx];
+}
+
+template <typename T>
+void launchBatchedTranspose2D(
+    T const* src, T* dst, int32_t outer, int32_t rows, int32_t cols, cudaStream_t stream)
+{
+    int64_t const total = static_cast<int64_t>(outer) * rows * cols;
+    if (total == 0)
+    {
+        return;
+    }
+    int32_t constexpr kBlockSize = 256;
+    int32_t const numBlocks = static_cast<int32_t>((total + kBlockSize - 1) / kBlockSize);
+    batchedTranspose2DKernel<T><<<numBlocks, kBlockSize, 0, stream>>>(src, dst, outer, rows, cols);
+}
+
+// Explicit instantiations.
+template void launchBatchedTranspose2D<float>(float const*, float*, int32_t, int32_t, int32_t, cudaStream_t);
+template void launchBatchedTranspose2D<half>(half const*, half*, int32_t, int32_t, int32_t, cudaStream_t);
+template void launchBatchedTranspose2D<__nv_bfloat16>(
+    __nv_bfloat16 const*, __nv_bfloat16*, int32_t, int32_t, int32_t, cudaStream_t);
+template void launchBatchedTranspose2D<int32_t>(
+    int32_t const*, int32_t*, int32_t, int32_t, int32_t, cudaStream_t);
+
+} // namespace nvinfer1::plugin

plugin/topkLastDimPlugin/transpose.h

@@ -0,0 +1,33 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 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_TOPK_LAST_DIM_TRANSPOSE_H
+#define TRT_TOPK_LAST_DIM_TRANSPOSE_H
+
+#include <cstdint>
+#include <cuda_runtime_api.h>
+
+namespace nvinfer1::plugin
+{
+
+//! Transpose the last two dimensions of a 3D tensor: [outer, rows, cols] -> [outer, cols, rows].
+template <typename T>
+void launchBatchedTranspose2D(T const* src, T* dst, int32_t outer, int32_t rows, int32_t cols, cudaStream_t stream);
+
+} // namespace nvinfer1::plugin
+
+#endif // TRT_TOPK_LAST_DIM_TRANSPOSE_H