Argmax Kernel

template<typename T>
void trt_edgellm::kernel::invokeRowwiseArgmax(

T const *input,

int32_t rows,

int32_t cols,

int32_t *outIndices,

cudaStream_t stream

)

Row-wise top-1 (argmax) over a [rows, cols] row-major tensor.

outIndices[r] = argmax_c input[r, c]. Ties resolve to the LOWEST column index, matching torch.argmax — required by callers that must reproduce a PyTorch greedy reference (RNN-T greedy decode, EAGLE target argmax).

One block per row; a warp-shuffle + shared-memory reduction inside the block. The comparison is done in float (T loads are widened), so the result is independent of reduction order and bit-stable across launch shapes.

Instantiated for T in {__half, float}.

Parameters:

• input – Device [rows, cols] row-major, dtype T.

• rows – Number of rows (blocks launched).

• cols – Row width (reduction extent).

• outIndices – Device [rows] int32, receives the argmax column per row.

• stream – CUDA stream.