[CUDA] Parallelize ArgMax/ArgMin over the reduction axis

[Tianlei Wu (tianleiwu)]

Description

arg_min_max_last_axis_kernel parallelizes only across rows, so a single thread walks the entire reduction axis. When the axis is long and there are few rows, the GPU is almost entirely idle.

This adds arg_min_max_last_axis_block_kernel, which assigns one block per row, strides the reduction axis across the block so the loads coalesce, and combines the per-thread candidates with a shared-memory tree reduction. The launcher picks it once the axis is at least a full block wide (256) and keeps the existing thread-per-row kernel for short axes, where that mapping is still the better one.

Motivation and Context

An ArgMax over a 129280-wide axis with a single row currently launches:

grid = (1, 1, 1)   block = (256, 1, 1)

with exactly one active thread, issuing 129280 dependent global loads — about 3.10 ms per call, roughly 0.17 GB/s of achieved bandwidth.

That shape is common at decode time, where a model takes an ArgMax over the vocabulary to pick the next token. It was found while profiling a speculative-decoding workload on 8×H200, where five such ArgMax calls per step accounted for 90.4% of the drafter's GPU kernel time.

Performance

Measured on H200 (sm90), one row, 129280-wide axis:

	per call
thread-per-row (before)	3098 µs
block-per-row (after)	46 µs
	67× faster

In the speculative-decoding workload that motivated the change, the affected model's kernel time dropped from 448.71 ms to 6.73 ms and end-to-end decode throughput improved by 52.6%, with unchanged output.

Semantics

Ties continue to resolve to the lowest index, matching ONNX's first-occurrence rule. Each thread seeds with element 0, which is always a valid candidate and can never displace a strictly better one, so rows shorter than the block need no separate guard. The CUDA EP already rejects select_last_index == 1, so only first-occurrence behavior has to be preserved.

Short axes (n < 256) keep the original kernel, since one thread per row is the better mapping when there are many rows and little to reduce.

Testing

• Compared the CUDA EP against the CPU EP over 396 configurations — float / MLFloat16 / double × ArgMax / ArgMin × keepdims 0/1 × random, tied, and all-equal inputs, across axis lengths spanning both sides of the 256 threshold (1, 2, 255, 256, 257, 1000, 129280) and row counts 1–300. All matched.
• Added ArgMax_float_long_axis_multiple_rows and ArgMin_float_long_axis_multiple_rows to reduction_ops_test.cc. They use a 300-wide axis — wider than the block and not a multiple of it — over 4 rows, covering the strided-loop remainder, the grid mapping across rows, a maximum in the first position, a maximum in the remainder region, duplicate extrema (must return the lower index), and a fully constant row.
• The existing ArgMax_float_first_index_random test already exercises this path (1 row, 65536-wide axis, duplicate +inf values expecting the lowest index) and continues to pass.

[Justin Chu (justinchuby)]

Superseded by #32092. Thanks for reviewing!