[CUDA] Parallelize ArgMax/ArgMin over wide last axes - #32092
Conversation
`arg_min_max_last_axis()` assigns one thread per row and scans the row serially. That is a good fit for many narrow rows, but it collapses when a model reduces a few very long rows: a `[1, 202048]` ArgMax over sampler logits is computed by a single CUDA thread and takes ~5 ms on an H200, leaving 99.9% of the device idle. The cost grows linearly with the reduction length, so it hurts every large-vocabulary classification and sampling graph, not one particular model. Add a cooperative kernel that spreads a row over many warps and, when the row is long enough, over many blocks. Partial (value, index) pairs are merged with warp shuffles, then through shared memory, and finally - for multi-block rows - through a small global buffer that the last block arriving for the row reduces. This is the same `__threadfence()` + per-row done counter structure that `reduce_matrix_columns()` already uses in this file, so the intermediate buffer follows the existing sizing/alignment conventions and is allocated by the caller through `AllocateScratchBuffer`. The launch geometry targets a single wave of threads over the device while keeping at least four elements per thread, which avoids both the idle-device case (few wide rows) and the oversubscription case (many rows). Rows narrower than 128 columns keep the existing one-thread-per-row kernel, which is still the better choice there, so no shape regresses. Semantics are unchanged. The reduction identity is -/+infinity (not the lowest/highest finite value) so a row containing -inf cannot lose against an empty partial result, ties resolve to the lowest index (`select_last_index == 0`; the CUDA EP already falls back to CPU when it is 1), and a leading NaN still yields index 0, matching the sequential kernel bit for bit for every input. Measured on H200 (sm90, CUDA 13.3), fp32, per call including the launch: | rows | cols | before | after | speedup | |------|--------|-----------|----------|---------| | 1 | 32000 | 695.6 us | 5.9 us | 118x | | 1 | 128256 | 3009.7 us | 6.1 us | 496x | | 1 | 202048 | 5056.5 us | 6.4 us | 793x | | 1 | 262144 | 6335.5 us | 6.3 us | 1005x | | 8 | 202048 | 5118.4 us | 8.7 us | 591x | | 32 | 202048 | 7357.7 us | 10.7 us | 688x | | 4096 | 128 | 18.8 us | 5.0 us | 3.7x | | 4096 | 4096 | 552.2 us | 20.7 us | 26.7x | | 65536| 1024 | 270.1 us | 67.7 us | 4.0x | | 4096 | 8 | 2.3 us | 2.2 us | 1.0x | | 65536| 64 | 18.9 us | 18.9 us | 1.0x | For `[1, 202048]` the launch is 197 blocks of 256 threads, 32 registers per thread, 100% theoretical occupancy, 3.8 us of kernel time plus a 4 byte memset, and 1595 bytes of scratch. The scratch is bounded by the device thread capacity, so it stays in the low kilobytes for any shape. Tests: adds arg reduction coverage to the CUDA reduction function tests (dispatch paths, dtypes, ties, NaN, infinities, buffer offsets, undersized buffer, CUDA graph capture/replay) and wide last-axis ArgMax/ArgMin operator tests, plus a disabled benchmark that prints the timing table above. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
|
Verification log for the numbers and claims in the description. H200 (sm90), CUDA 13.3, cuDNN 9.24, CUDA reduction unit tests, including CUDA graph capture and replay
The first test to run in the module absorbs a one time ~23 s CUDA module load; the arg reduction tests themselves are 0-511 ms. Operator level testsBenchmark, raw output of the added benchmark on both kernels
The "before" column is the same benchmark with the cooperative path disabled, so both columns come from the same binary, the same harness and the same data. Kernel profile and launch geometryNsight Systems, So 1 kernel launch plus one 4 byte memset per call, 4.69 us of GPU time in total against ~5056 us before. Launch geometry and occupancy ( Increasing the block count past the heuristic makes it slower rather than faster, which is why the geometry targets a single wave: |
The thread level scan advanced its position with `id += MAX_NUM_ELEMENTS_PER_THREAD * num_threads_in_grid_row` and formed each load offset as `id + i * num_threads_in_grid_row`, both in `int`. A row is admitted up to `INT_MAX` columns and a grid row holds up to 256 blocks x 256 threads, so the step reaches 262144 and those sums overflow to negative values once `id` approaches `INT_MAX`. That reads far outside the tensor and, because a negative position still satisfies `id < num_cols`, the loop never terminates. Reproduced on an H200 with a `[1, INT_MAX]` fp16 row: 12884901888 negative positions computed and every one of the 65536 threads still looping when a watchdog stopped them. Write the scan with differences instead of sums. `remaining = num_cols - id` and `delta = i * num_threads_in_grid_row` are both smaller than `num_cols`, and a position is only formed as `id + delta` after `delta < remaining` has been checked, so every value stays inside `[0, num_cols)` and no sum can overflow. The guard `delta < remaining` is exactly the old `id + delta < num_cols`, and the recorded index is unchanged, so results are identical for every input. int64 positions were measured as the alternative. They are correct as well but cost 4% to 38% on the shapes this PR targets (for example [65536, 256] 28.0 us -> 38.7 us, [32, 202048] 10.7 us -> 13.2 us) because the loads lose 32 bit addressing, while the difference based form matches the previous throughput ([1, 202048] 6.4 us -> 6.2 us) and needs 32 registers instead of 40 for double. No supported shape is dropped either way: the dispatch range is unchanged. Tests: a `[1, INT_MAX]` fp16 case (skipped when the device cannot hold the ~4 GiB allocation), widths straddling the widest per iteration scan step, and buffer sizing at `INT_MAX` widths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
|
Thanks, this is a real bug and it is now fixed in 0d85ddf. Details and evidence below. I am leaving the review thread open for you to close. The bugThe thread level scan advanced with Reproduced on H200 with a So every thread would have spun forever, and 1.3e10 reads would have landed outside the tensor. The fixThe scan is now expressed with differences instead of sums: const int scan_step = MAX_NUM_ELEMENTS_PER_THREAD * num_threads_in_grid_row;
for (int id = tid_in_grid_row; id < num_cols;) {
const int remaining = num_cols - id; // 1 <= remaining <= num_cols
...
const int delta = i * num_threads_in_grid_row;
if (delta < remaining) { v[i] = row_data[id + delta]; } // id + delta < num_cols
...
if (remaining <= scan_step) break;
id += scan_step; // scan_step < remaining, so id stays < num_cols
}Every quantity is provably inside Why not int64 positionsI implemented and measured that variant too, since it was the suggested fix. It is correct, but the loads lose 32 bit addressing and it costs 4% to 38% on exactly the shapes this PR is about, while the difference based form matches the previous throughput. Register use is also equal or better (double: 32 vs 40 registers, so it keeps a higher occupancy).
Registers,
No supported shape is narrowed: the dispatch conditions in Primary shape, no regression
New tests
Verification after the fixH200 (sm90), CUDA 13.3, Release, |
There was a problem hiding this comment.
Pull request overview
Parallelizes CUDA ArgMax/ArgMin reductions over wide last axes while preserving existing semantics.
Changes:
- Adds cooperative warp/block reduction with scratch-buffer support.
- Integrates shape-aware dispatch into CUDA reduction operators.
- Adds extensive kernel, operator, edge-case, and performance tests.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
reduction_functions.cu |
Implements cooperative CUDA reduction. |
reduction_functions.h |
Exposes buffer sizing and updated API. |
reduction_ops.cc |
Allocates scratch space and dispatches kernels. |
reduction_functions_test.cc |
Adds CUDA kernel and benchmark coverage. |
reduction_ops_test.cc |
Adds end-to-end wide-axis tests. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Design decision: no ORT custom-op registration. Three implementations remain, and none of them needs a custom domain: 1. the native EP device-sampler kernel (#976, main-targeted), 2. generic upstream ORT CUDA ArgMax (microsoft/onnxruntime#32092), 3. the portable ONNX tiled rewrite here, for stock ORT versions that predate #32092. Removes the custom op and everything that existed only to serve it: the fused_argmax module and its session test, the CUDA-only custom-op domain registration in Session::new, the prost dev-dependency that only that test needed, the WideArgReduceLowering Fused variant with its fusable, fuse and element_types helpers and FUSABLE_ELEMENT_TYPES, the ArgReduceRewrites fused counter, the FUSED_DOMAIN and FUSED_OP constants and their re-export, the opset-import injection for the custom domain, the availability probe in islands, and the fused benchmark and test modules. Direct and Tiled are untouched, and so is the version gate that chooses between them, so a future ORT containing #32092 still gets the node as authored. Coverage that was only reachable through the fused test modules is preserved as a direct module rather than deleted. On an H200 against stock ORT 1.28 the gate still reports direct-capable false, and the surviving sweep measures direct 35538.28 us against tiled 155.34 us over 1 and 10 nodes, which is why the portable fallback stays. The lowering parity test passes, so Direct and Tiled still agree on the selected token. engine 324 lib tests, ort 135 lib tests and the 3 ignored CUDA argmax tests pass; cargo check clean for onnx-genai-ort with and without the cuda feature and for onnx-genai-engine. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chuby <justinchuby@users.noreply.github.com>
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
Reviewed the cooperative reduction, dispatch, scratch-buffer layout, and focused tests. The multi-block completion protocol is correctly ordered (partial write, device-wide fence, per-row atomic completion), tie-breaking preserves the lowest index, leading-NaN behavior matches the existing serial path, and the difference-based scan covers the admitted INT_MAX width without signed overflow. The graph replay, alignment, invalid-buffer, special-value, and width-boundary tests give this good coverage.
Compared with #31694, this PR should merge and #31694 should be closed as superseded. #31694's one-256-thread-block-per-row design is a useful improvement over the serial scan, but it still occupies only one SM for the motivating [1, vocab_size] workload. This PR adaptively uses multiple blocks per row while retaining the serial path for narrow rows, and both PRs modify the same launcher and overlapping tests. If benchmarks on other GPU generations later identify a middle-width range where one fixed block wins, that dispatch choice can be folded into this implementation rather than maintaining two competing paths.
The existing benchmark-command thread is documentation-only and does not block approval.
Description
arg_min_max_last_axis()inreduction_functions.cuassigns one thread per row and scans the row serially. That is a reasonable design for many narrow rows, but it degenerates badly when a graph reduces a small number of very long rows: the reduction length becomes a serial dependency chain in a single thread, so the cost grows linearly with the reduced axis and the rest of the GPU stays idle.A
[1, 202048]ArgMax(last-axis classification / sampling over a large vocabulary) is therefore computed by exactly one CUDA thread and takes ~5 ms on an H200. The same shape costs microseconds once it is parallelized. This affects any model that reduces a wide last axis with few rows, and it also leaves a lot on the table for the medium cases (thousands of rows of a few hundred to a few thousand columns).This PR adds a cooperative reduction path and keeps the existing serial kernel for the shapes where it is genuinely better.
Approach
__threadfence()+ per-row done-counter structure thatreduce_matrix_columns()/reduce_all()already use in the same file, so the intermediate buffer follows the existing sizing and alignment conventions and is allocated by the caller withAllocateScratchBuffer(same as theApplicableMatrixReduction::Columnspath).Semantics
The new kernel is bit-identical to the previous kernel for every input, including:
select_last_index == 0. The CUDA EP already falls back to CPU forselect_last_index == 1(ArgMaxOrArgMinNeedFallbackToCPU), so this is the only case to support.-inf/+inf, not the lowest/highest finite value. This matters: with a finite identity, a row of[-inf, ..., lowest_finite, ...]would drop the real maximum. Rows that are entirely-inf,+infor NaN return index 0, as before.axis,keepdims, negative axes, opset variants and the int64 output are untouched, they are handled above this function.MLFloat16,float,doublefor ArgMax/ArgMin) and the fallbacks for non-last-axis reductions,n/mbeyondintrange and empty reductions are unchanged.Results
H200 (sm90), CUDA 13.3, fp32, per call including the launch, measured with the benchmark added in this PR (
ReductionFunctionsTest.DISABLED_ArgMinMaxLastAxisPerf), "before" measured with the same harness on the serial kernel:Narrow rows are unchanged because they keep the existing kernel.
halfanddoublebehave the same way ([1, 202048]: 2432 us -> 6.1 us forhalf, 6405 us -> 6.5 us fordouble).Profile of
[1, 202048]fp32 (Nsight Systems, occupancy fromcudaOccupancyMaxActiveBlocksPerMultiprocessor):For reference on the same GPU and shape,
cub::DeviceReduce::ArgMaxneeds 2 kernels, 2.74 us of GPU time and 42495 bytes of temporary storage. The kernel here is a segmented (per row) reduction that also has to serve the many-rows shapes, uses ~26x less scratch and one launch instead of two. Scratch is bounded by the device thread capacity (a multi-block launch only happens whenrows < capacity / 256), so it stays in the low kilobytes for any shape.Testing
onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc: newArgMinMaxLastAxis*cases covering every dispatch path (narrow / single block / multi block / more rows than grid rows),half/float/double, ArgMax and ArgMin, ties, NaN (leading, middle, all),+/-inf, all-equal rows, finite extremes, unaligned and undersized intermediate buffers, a dirtied intermediate buffer, and CUDA graph capture + repeated replay.onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc: operator level ArgMax/ArgMin tests with a wide last axis ({3, 40000}and{2, 202048}, withkeepdimson and off and a negative axis) so the CUDA EP scratch allocation path is covered end to end.onnxruntime_provider_test --gtest_filter='*ArgMax*:*ArgMin*': 32/32 pass.onnxruntime_provider_test --gtest_filter='*Reduce*': 376/376 pass.CUDA_EP_Unittest.All(CUDA EP internal tests): 76/76 pass.compute-sanitizer(memcheck,racecheck,synccheck) clean on the multi-block and row-loop paths.{+/-inf, NaN, +/-0.0, ...}alphabet, repeated with several simulated device capacities: no mismatches.Review follow-up
idinint. Since a row is admitted up toINT_MAXcolumns and the step reaches4 * 65536 = 262144, that overflowed to negative positions for the widest rows, reading out of bounds and never terminating. Fixed by expressing the scan with differences (remaining = num_cols - id,delta = i * threads, positions only formed afterdelta < remaining), which keeps every value inside[0, num_cols). int64 positions were measured as the alternative and were 4-38% slower on the targeted shapes, so the difference based form was kept; no shape is dropped either way. New tests:[1, INT_MAX]fp16 (memory gated), widths straddling the widest scan step, and buffer sizing atINT_MAX.Possible follow-up
The thread level scan uses scalar loads. Vectorized loads would shave roughly another microsecond off the very wide single row case; it was left out here to keep the alignment handling out of a generic kernel that has to serve
half,floatanddouble.