[CUDA Backend] Optimize TopKV2 OP for CUDA backend

[rainyl]

Motivation

I am benchmarking RT-DETRv4 and YOLO26 models using MNN and find that they are much slower than other models like YOLO11, timeProfile tells that TopKV2 for CUDA is much slower than CPU, which is strange, so I checked source/backend/cuda/execution/TopKV2Execution.cu using LLM and found insertion sort is used in TopKInThread, which is expensive as it's time complexity is O(K), especially for large K like 300 used in RT-DETRv4 and YOLO26.

What's changed

Replace the insertion sort with heap sort to reduce the time complexity to O(log K), and use vectorized load if possible to speed up the TopKV2 OP on CUDA backend.

Comparision

run_test.out op/TopKV2

Before

./build/run_test.out op/TopKV2 2
type=2 in main, 37 
thread=1 in main, 50 
precision=1 in main, 51 
memory=0 in main, 56 
CPU Group: [ 20  21  31  13  23  1  15  25  3  17  27  5  19  29  7  10  11  30  9  12  22  0  14  24  2  16  26  4  18  28  6  8 ], 600000 - 5752000
The device supports: i8sdot:0, fp16:0, i8mm: 0, sve2: 0, sme2: 0
        running op/TopKV2.
0 s 19 ms
op/TopKV2 cost time: 74.950 ms
√√√ all <op/TopKV2> tests passed.
TEST_NAME_UNIT: 单元测试
TEST_CASE_AMOUNT_UNIT: {"blocked":0,"failed":0,"passed":1,"skipped":0}
TEST_CASE={"name":"单元测试","failed":0,"passed":1}

After

./build/run_test.out op/TopKV2 2
type=2 in main, 37 
thread=1 in main, 50 
precision=1 in main, 51 
memory=0 in main, 56 
CPU Group: [ 20  21  31  13  23  1  15  25  3  17  27  5  19  29  7  10  11  30  9  12  22  0  14  24  2  16  26  4  18  28  6  8 ], 600000 - 5752000
The device supports: i8sdot:0, fp16:0, i8mm: 0, sve2: 0, sme2: 0
        running op/TopKV2.
0 s 3 ms
op/TopKV2 cost time: 60.398 ms
√√√ all <op/TopKV2> tests passed.
TEST_NAME_UNIT: 单元测试
TEST_CASE_AMOUNT_UNIT: {"blocked":0,"failed":0,"passed":1,"skipped":0}
TEST_CASE={"name":"单元测试","failed":0,"passed":1}

Note that the time cost print by printTimeCost:

• before: 0 s 19 ms
• after: 0 s 3 ms

Which is a huge improvement.

RT-DETRv4 benchmark

Before

// ./build/timeProfile.out my_models/rt-detrv4-s-0_q0.mnn 100 2 1x3x640x640 4 2
Sort by time cost !
Node Type            	Avg(ms)   	%         	Called times 	Flops Rate 	
Cast                 	0.006860  	0.020853  	1.000000     	0.000000   	
Interp               	0.018980  	0.057695  	2.000000     	0.013797   	
Pooling              	0.020270  	0.061616  	1.000000     	0.044151   	
ReLU6                	0.041860  	0.127245  	9.000000     	0.003129   	
Reduction            	0.055810  	0.169651  	8.000000     	0.002150   	
Softmax              	0.072000  	0.218864  	11.000000    	0.024824   	
LayerNorm            	0.084810  	0.257803  	12.000000    	0.020523   	
ConvolutionDepthwise 	0.110050  	0.334526  	14.000000    	0.233287   	
GridSample           	0.146469  	0.445231  	9.000000     	0.074504   	
While                	0.171380  	0.520955  	25.000000    	1.507162   	
UnaryOp              	0.866189  	2.633011  	65.000000    	0.108062   	
BinaryOp             	0.968978  	2.945464  	183.000000   	0.259645   	
Convolution          	1.528610  	4.646615  	139.000000   	97.330925  	
Raster               	1.840351  	5.594233  	220.000000   	0.376477   	
TopKV2               	26.968828 	81.978882 	3.000000     	0.000036   	
total time : 32.897289 ms, total mflops : 14156.047852 
main, 171, cost time: 3354.350098 ms

After

Sort by time cost !
Node Type            	Avg(ms)  	%         	Called times 	Flops Rate 	
Cast                 	0.005940 	0.066846  	1.000000     	0.000000   	
Interp               	0.018480 	0.207967  	2.000000     	0.013797   	
Pooling              	0.019940 	0.224397  	1.000000     	0.044151   	
ReLU6                	0.042350 	0.476593  	9.000000     	0.003129   	
Reduction            	0.055290 	0.622217  	8.000000     	0.002150   	
Softmax              	0.072620 	0.817239  	11.000000    	0.024824   	
LayerNorm            	0.085300 	0.959936  	12.000000    	0.020523   	
ConvolutionDepthwise 	0.109070 	1.227425  	14.000000    	0.233287   	
GridSample           	0.148299 	1.668899  	9.000000     	0.074504   	
While                	0.170520 	1.918965  	25.000000    	1.507162   	
UnaryOp              	0.874023 	9.835900  	65.000000    	0.108062   	
BinaryOp             	0.966012 	10.871109 	183.000000   	0.259645   	
Convolution          	1.543254 	17.367165 	139.000000   	97.330925  	
Raster               	1.853716 	20.860973 	220.000000   	0.376477   	
TopKV2               	2.920069 	32.861279 	3.000000     	0.000036   	
total time : 8.886047 ms, total mflops : 14156.047852 
main, 171, cost time: 949.072998 ms

For TopKV2 OP, it speeds up 26.968828/2.920069=9.23X faster.

[wangzhaode]

Code Review

1. Duplicated sift-down logic (Maintainability)

The sift-down heap logic is copy-pasted three times: in the unrolled loop, in the remainder loop, and in the heap-sort phase. This is a significant amount of duplicated device code. Consider extracting a __device__ inline helper function like SiftDown(valuesThread, indicesThread, parent, heapSize, descendFlag) to reduce duplication and maintenance burden.

2. Unrolled loop may not achieve vectorized loads as claimed (Comments accuracy)

The comments suggest the unrolling helps the compiler generate vectorized loads (LD.E.128). However, the access pattern uses a stride of gridDim.x * blockDim.x between consecutive elements in data[0..3], so these are not contiguous memory accesses. Vectorized loads like float4 require contiguous addresses. The unrolling still helps with instruction-level parallelism (ILP) and latency hiding, but the comments about vectorized loads are misleading. Consider correcting the comments to mention ILP benefits instead.

Reviewed by claude-opus-4-6