Skip to content

Add INT8, INT16, and UINT8 type support for CUDA TopK operator #27862

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tianleiwu merged 3 commits into from
Apr 19, 2026
Merged

Add INT8, INT16, and UINT8 type support for CUDA TopK operator #27862

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f1a7599
Add INT8, INT16, and UINT8 type support for CUDA TopK operator
elwhyjay Mar 26, 2026
88a9f34
Restrict int8/int16/uint8 TopK CUDA kernels to opset 11+
elwhyjay Mar 27, 2026
35d7a76
Update docs/OperatorKernels.md for CUDA TopK int8/int16/uint8 support
elwhyjay Mar 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions docs/OperatorKernels.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -996,8 +996,8 @@ Do not modify directly.*
|||1+|**T** = tensor(double), tensor(float), tensor(float16)|
|Tile|*in* input:**T**<br> *in* repeats:**T1**<br> *out* output:**T**<br><br>or<br><br>*in* input:**T**<br> *in* tiles:**T**<br> *in* axis:**T**<br> *out* output:**T**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64)<br/> **T1** = tensor(int64)|
|||[6, 12]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64)<br/> **T1** = tensor(int64)|
|TopK|*in* X:**T**<br> *in* K:**tensor(int64)**<br> *out* Values:**T**<br> *out* Indices:**I**<br><br>or<br><br>*in* X:**T**<br> *out* Values:**T**<br> *out* Indices:**I**|24+|**I** = tensor(int64)<br/> **T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64)|
|||[11, 23]|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64)|
|TopK|*in* X:**T**<br> *in* K:**tensor(int64)**<br> *out* Values:**T**<br> *out* Indices:**I**<br><br>or<br><br>*in* X:**T**<br> *out* Values:**T**<br> *out* Indices:**I**|24+|**I** = tensor(int64)<br/> **T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)|
|||[11, 23]|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint8)|
|||10|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64)|
|||[1, 9]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64)|
|Transpose|*in* data:**T**<br> *out* transposed:**T**|23+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
Expand Down
11 changes: 10 additions & 1 deletion onnxruntime/core/providers/cuda/math/topk.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,10 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()})
DataTypeImpl::GetTensorType<int64_t>(),
DataTypeImpl::GetTensorType<int8_t>(),
Comment thread
tianleiwu marked this conversation as resolved.
DataTypeImpl::GetTensorType<int16_t>(),
DataTypeImpl::GetTensorType<uint8_t>()})
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>()),
TopK<true>);

Expand All @@ -62,6 +65,9 @@ ONNX_OPERATOR_KERNEL_EX(
DataTypeImpl::GetTensorType<double>(),
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>(),
DataTypeImpl::GetTensorType<int8_t>(),
DataTypeImpl::GetTensorType<int16_t>(),
DataTypeImpl::GetTensorType<uint8_t>(),
DataTypeImpl::GetTensorType<BFloat16>()})
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>()),
TopK<true>);
Expand Down Expand Up @@ -137,6 +143,9 @@ Status TopK<inputk>::ComputeInternal(OpKernelContext* ctx) const {

if (IS_PRIM_TYPE(int32_t)) return TOPKIMPL(int32_t);
if (IS_PRIM_TYPE(int64_t)) return TOPKIMPL(int64_t);
if (IS_PRIM_TYPE(int8_t)) return TOPKIMPL(int8_t);
if (IS_PRIM_TYPE(int16_t)) return TOPKIMPL(int16_t);
if (IS_PRIM_TYPE(uint8_t)) return TOPKIMPL(uint8_t);
if (IS_PRIM_TYPE(MLFloat16)) return TOPKIMPL(MLFloat16);
if (IS_PRIM_TYPE(float)) return TOPKIMPL(float);
if (IS_PRIM_TYPE(double)) return TOPKIMPL(double);
Expand Down
5 changes: 5 additions & 0 deletions onnxruntime/core/providers/cuda/math/topk_impl_i16.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#define TOPK_IMPL_TYPE int16_t
#include "topk_impl.cuh"

Check warning on line 5 in onnxruntime/core/providers/cuda/math/topk_impl_i16.cu

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Include the directory when naming header files  [build/include_subdir] [4]

Raw Output:
onnxruntime/core/providers/cuda/math/topk_impl_i16.cu:5:  Include the directory when naming header files  [build/include_subdir] [4]
5 changes: 5 additions & 0 deletions onnxruntime/core/providers/cuda/math/topk_impl_i8.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#define TOPK_IMPL_TYPE int8_t
#include "topk_impl.cuh"

Check warning on line 5 in onnxruntime/core/providers/cuda/math/topk_impl_i8.cu

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Include the directory when naming header files  [build/include_subdir] [4]

Raw Output:
onnxruntime/core/providers/cuda/math/topk_impl_i8.cu:5:  Include the directory when naming header files  [build/include_subdir] [4]
5 changes: 5 additions & 0 deletions onnxruntime/core/providers/cuda/math/topk_impl_u8.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#define TOPK_IMPL_TYPE uint8_t
#include "topk_impl.cuh"

Check warning on line 5 in onnxruntime/core/providers/cuda/math/topk_impl_u8.cu

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Include the directory when naming header files  [build/include_subdir] [4]

Raw Output:
onnxruntime/core/providers/cuda/math/topk_impl_u8.cu:5:  Include the directory when naming header files  [build/include_subdir] [4]
Loading