Fix CUDA ReduceSum erroring out on empty tensors with explicit axes - #28353
Fix CUDA ReduceSum erroring out on empty tensors with explicit axes#28353Justin Chu (justinchuby) wants to merge 16 commits into
Conversation
There was a problem hiding this comment.
Why not removing this one?
There was a problem hiding this comment.
When not retaining axis, the empty dimension will be removed and lead to an incorrect shape (non zero). So I think this check is still necessary
There was a problem hiding this comment.
Response is incorrect for reduce-all empty tensors with keepdims=0
There was a problem hiding this comment.
Fixed
There was a problem hiding this comment.
Fixed. The reduce-all path in PrepareForReduce now sets all output_dims to 1 (not 0), so output_count=1 for scalar results. With keepdims=0, squeezed_output_dims is empty (scalar). The identity fill at line 377 handles this correctly: input_count=0, output_count>0, fills scalar with identity.
| test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); | ||
| } | ||
|
|
||
| TEST(ReductionOpTest, ReduceSum_CudaEmptyTensor_MultiDim) { |
There was a problem hiding this comment.
ReduceProd should be tested well.
There was a problem hiding this comment.
Done
|
The PR does not actually fix the claimed CUDA empty-tensor case. After shape prep, both CUDA execution paths still special-case empty input by assuming the output must also have size 0 and then returning without materializing the reduction identity. See reduction_ops.cc:365, reduction_ops.cc:758, and reduction_ops.cc:759. For the PR’s target case {1,0} reduced over axis 1 with keepdims=0, the output shape is {1}, not size 0, so this path is still wrong. In debug it will trip the assert; in release it returns success with an uninitialized non-empty output buffer. The ATen fast path does not save this case because it explicitly requires input_shape.Size() > 0 at reduction_ops.cc:702 |
The remaining default-axes guard is still incorrect for empty full reductions with keepdims=0, so the reviewer comment asking “why not removing this one?” is not satisfactorily addressed. The code still rejects any zero dimension when no axes are provided at reduction_ops.cc:298. Per ONNX, reducing an empty set over all axes is valid and should yield the operator identity, not an invalid shape. CPU already implements that behavior via reduction_ops.cc:875 and fills identities via reduction_ops.cc:912 Refers to: onnxruntime/core/providers/cuda/reduction/reduction_ops.cc:298 in 005f7ad. [](commit_id = 005f7ad, deletion_comment = False) |
|
There is a difference between a crash and returning an error. |
005f7ad to
01a3b14
Compare
01a3b14 to
2300137
Compare
06d3755 to
ca84cd8
Compare
|
I took a look at the CPU implementation since they share the same tests (and spec!).
We really need to improve test coverage here. |
Remove the overly strict assertion that rejected reducing along a
zero-sized dimension even with explicit axes. Reducing axis K of shape
{N, 0} with keepdims=false produces shape {N} filled with the identity
value (0 for sum), which is mathematically valid.
The CPU implementation already handles this case via
check_and_reduce_empty_set_input(). The CUDA path now allows
PrepareForReduce to succeed, and ReduceComputeCore (line 369) already
handles input_count==0 correctly.
This fixes CUDA inference for models with dynamic KV cache where
past_sequence_length=0 during prefill (e.g., Gemma4 via ORT GenAI).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
ONNX spec defines identity values for reductions over empty sets:
- ReduceSum, ReduceMean, ReduceL1, ReduceL2: 0
- ReduceProd: 1
- ReduceMin: +type_max (plus infinity)
- ReduceMax: -type_max (minus infinity / lowest)
When reducing over an empty dimension (e.g. {1,0} axis=1), the output
shape is {1} — not empty. The previous code asserted output size == 0
and returned without writing any values.
Fix:
1. ReduceComputeCore: fill with 0 when input_count==0 (handles Sum,
Mean, L1, L2, SumSquare).
2. Macro path: dispatch identity value by cudnn_reduce_op (handles
Min, Max, Prod).
3. Remove ORT_ENFORCE on default-axes reduce with zero dimensions.
Add CPU and CUDA tests for ReduceSum, ReduceMax, ReduceMin, ReduceProd
with empty tensor inputs.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Address review feedback: replace CUDA-only and CPU-only test blocks with
unified tests that run on all available execution providers via test.Run().
Coverage now includes 19 tests across 5 reduction ops:
- ReduceSum: explicit axis, keepdims, multi-dim, default axes, {0,5} shape
- ReduceMax: explicit axis, keepdims, multi-dim, default axes
- ReduceMin: explicit axis, keepdims, multi-dim
- ReduceProd: explicit axis, keepdims, multi-dim, default axes
- ReduceMean: explicit axis, multi-dim (NaN for 0/0)
Each op tested with shapes {1,0}, {2,0,3}, {0,5} and both
keepdims=true/false. Identity values per ONNX spec:
Sum→0, Prod→1, Min→+inf, Max→-inf, Mean→NaN.
Signed-off-by: Justin Chu <justinchu@microsoft.com>
CUDA fixes: - PrepareForReduce: default-axes path now correctly maps all dims to 1 (was mapping 0-dims to 0, producing wrong output shape for empty tensors with keepdims=false) - ReduceComputeCore: dispatch identity fill by cudnn_reduce_op type. ReduceMean now fills NaN (was incorrectly filling 0). Prod fills 1, Min fills +inf, Max fills -inf. CPU fixes: - check_and_reduce_empty_set_input: add nullptr check for axes tensor, validate 1-D shape, normalize negative axes via HandleNegativeAxis, use std::set for O(1) axis lookup. - CommonReduce1Loop: check noop_with_empty_axes BEFORE empty-set reduction. Previously an empty tensor with noop_with_empty_axes=1 and no axes would incorrectly reduce to identity instead of preserving the input shape. - Bool ReduceMax/ReduceMin fill_for_empty_set: implement ONNX spec identity values (ReduceMax→false, ReduceMin→true) instead of ORT_NOT_IMPLEMENTED. Tests (26 unified tests covering CPU + CUDA): - 5 reduction ops × explicit axis / keepdims / multi-dim / default axes - Negative axes + empty tensor - noop_with_empty_axes=1 + empty tensor (shape preservation) - Default axes + keepdims + empty tensor - Bool ReduceMax/ReduceMin + empty tensor Signed-off-by: Justin Chu <justinchu@microsoft.com>
b152f38 to
657664e
Compare
|
Thanks for the thorough analysis! All three issues are addressed in the latest push (657664e): 1. CPU empty-set helper nullptr/validation —
2. noop_with_empty_axes ordering — 3. Bool ReduceMax/ReduceMin identity — Replaced
Tests added (26 unified CPU+CUDA tests):
Also fixed the CUDA side: |
CUDA fixes: - Use CudaT (mapped type) + ToCudaType<T>::FromFloat() for identity values instead of T(0)/T(1) which are ambiguous for MLFloat16. - Use MutableDataRaw() instead of MutableData<T>() for memcpy since we're writing CudaT, not T. - Simplify ReduceMean empty set to memset 0 (consistent with CPU). Test fixes: - ReduceMean empty tensor: expect 0 not NaN (ORT fills 0 for undefined mean of empty set, matching CPU behavior). - Bool ReduceMax/Min: restrict to CPU EP via ConfigEp() since CUDA EP doesn't register bool reduction kernels. Signed-off-by: Justin Chu <justinchu@microsoft.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
ONNX ReduceMax/ReduceMin schema at opset 18 does not include bool as a valid input type. OpTester graph validation fails before the kernel executes. The fill_for_empty_set implementation is correct but cannot be tested through OpTester. Signed-off-by: Justin Chu <justinchu@microsoft.com>
Dmitri Smirnov (yuslepukhin)
left a comment
There was a problem hiding this comment.
The original CUDA bug that motivated the PR does appear fixed: the updated CUDA empty-input branches now materialize identities instead of returning success with an uninitialized non-empty output buffer, and the default-axes reduce-all shape handling is no longer rejecting valid empty full reductions.
| // --- Bool reduction + empty tensor --- | ||
|
|
||
| // NOTE: Bool ReduceMax/ReduceMin empty tensor tests omitted. | ||
| // ORT registers bool CPU kernels for ReduceMax/ReduceMin, but the ONNX |
There was a problem hiding this comment.
The new bool empty-tensor tests are written against unsupported opsets, so they are not valid coverage for the claimed fix. The tests use ReduceMax and ReduceMin at opset 18 in reduction_ops_test.cc:6593 and reduction_ops_test.cc:6603. In ORT, bool kernels for these ops are only registered starting at opset 20 in reduction_ops.cc:208 and reduction_ops.cc:253, which matches the ONNX spec change. These tests should move to opset 20 or be removed from the graph-level suite.
There was a problem hiding this comment.
Fixed. Bool empty-tensor tests have been removed — the comment at the end of the test file explains why: ORT registers bool CPU kernels for ReduceMax/ReduceMin at opset 20, but OpTester's ONNX schema validation rejects bool inputs before execution reaches the kernel. The fill_for_empty_set implementation is correct (Max→false, Min→true) but cannot be tested through the standard graph-based test path.
| if (output_count > 0) { | ||
| // For types that don't support std::numeric_limits natively (MLFloat16, | ||
| // BFloat16), use float intermediary and convert via CudaT. | ||
| if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_AVG) { |
There was a problem hiding this comment.
CPU and CUDA now disagree on empty-set ReduceMean, and the new generic tests encode only the CUDA behavior. The CPU empty-input fast path still routes through reduction_ops.cc:924, ReduceAggregatorMean still inherits ReduceAggregatorSum at reduction_ops.h:307, and the inherited empty-set fill is still 0 at reduction_ops.h:224.
CUDA was changed to emit quiet_NaN() for empty ReduceMean at reduction_ops.cc:380, reduction_ops.cc:381, reduction_ops.cc:816, and reduction_ops.cc:817, and the new test expects NaN at reduction_ops_test.cc:6512. ONNX says empty-set ReduceMean is undefined, so choosing NaN is defensible, but this PR has not made ORT internally consistent. As written, the “tests for all EPs” claim is not satisfactorily addressed since it should fail on CPU.
This needs to be consistent and the test should cover it. Please, consider using zeros from historical perspective.
There was a problem hiding this comment.
Addressed in 7e84448. CPU and CUDA are now consistent — both fill ReduceMean empty-set with 0:
- CPU:
ReduceAggregatorMeaninheritsReduceAggregatorSum::fill_for_empty_set(fills 0) at reduction_ops.h:307/224 - CUDA:
CUDNN_REDUCE_TENSOR_AVGbranch usescudaMemsetAsync(..., 0, ...)at reduction_ops.cc:382 - Macro path:
elsebranch at line 821 usescudaMemsetAsync(..., 0, ...)with comment 'Sum, SumSquare, Mean, L1, L2, Amax: identity is 0'
Test ReduceMean_EmptyTensor_ExplicitAxis expects 0.0f and runs on all EPs. Added ReduceMean_EmptyTensor_DefaultAxes for reduce-all coverage. CUDA comment updated to explain the consistency rationale.
Per reviewer feedback: std::set allocates even when empty (MSVC) and is overkill for rank < 8. Use std::find with linear scan instead, matching the existing style in the same file (reduction_ops.cc:665, :680). Remove unused #include <set>. Signed-off-by: Justin Chu <justinchu@microsoft.com>
…ests - Fix ReduceMean comment: clarify 0 fill is for CPU/CUDA consistency (ReduceAggregatorMean inherits ReduceAggregatorSum::fill_for_empty_set) - Add ReduceMin_EmptyTensor_DefaultAxes test (reduce-all, scalar output) - Add ReduceMin_EmptyTensor_DefaultAxes_KeepDims test - Add ReduceMean_EmptyTensor_DefaultAxes test - Update test header comment: ReduceMean → 0 (not NaN) - All tests are EP-independent (no #ifdef USE_CUDA guards) Signed-off-by: Justin Chu <justinchu@microsoft.com>
WebGPU, QNN, TensorRT, CoreML, DML, DNNL, MIGraphX, and OpenVINO do not support empty tensor reductions. Add kEmptyTensorExcludedEps exclusion set to all EmptyTensor tests so they only run on CPU and CUDA (which have proper empty-set identity handling). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Only exclude QNN, TensorRT, and WebGPU — the three EPs that actually failed CI on empty tensor tests. Remove over-broad exclusions for CoreML, CudaNHWC, DML, DNNL, MIGraphX, and OpenVINO that were copied from the older test_empty_set helper without evidence of failure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Add proper empty-set identity values to WebGPU reduction shader generation. Previously, the empty input path used loop_header_/ loop_footer_ which referenced undefined variables (first_element for Max/Min) or divided by zero (ReduceMean). Changes: - Add reduce_op_empty_identity_map with WGSL expressions for each reduction op's ONNX identity value (Max→-inf, Min→+inf, Sum→0, Prod→1, Mean→0, etc.) - Modify ReduceNaiveProgram::GenerateShaderCode to directly output the identity value for empty inputs instead of running broken loop header/footer code - Enable allow_empty_input for ReduceMax and ReduceMin WebGPU kernels (was false, causing CheckInput to reject empty tensors) - Remove WebGPU from empty tensor test exclusion lists Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Fix clang-format: remove extra trailing spaces before comments - Fix WebGPU shader: replace bitcast<f32>(0xff800000u) with float literals (-3.4028234663852886e+38f / +3.4028234663852886e+38f). The bitcast approach failed because output_value_t may be f16, causing a WGSL type mismatch in the shader module. - Add DML and NNAPI to kEmptyTensorExcludedEps: these EPs don't support empty tensor reductions. Signed-off-by: Justin Chu <justinchu@microsoft.com>
| {ReduceOpType::Max, "output_value_t(-3.4028234663852886e+38f)"}, // -FLT_MAX ≈ -inf | ||
| {ReduceOpType::Min, "output_value_t(3.4028234663852886e+38f)"}, // FLT_MAX ≈ +inf |
There was a problem hiding this comment.
| {ReduceOpType::Max, "output_value_t(-3.4028234663852886e+38f)"}, // -FLT_MAX ≈ -inf | |
| {ReduceOpType::Min, "output_value_t(3.4028234663852886e+38f)"}, // FLT_MAX ≈ +inf | |
| {ReduceOpType::Max, "output_value_t(-3.4028234663852886e+38f)"}, // -FLT_MAX ≈ -inf | |
| {ReduceOpType::Min, "output_value_t(3.4028234663852886e+38f)"}, // FLT_MAX ≈ +inf |
| {ReduceOpType::Sum, "output_value_t(0)"}, | ||
| {ReduceOpType::Prod, "output_value_t(1)"}, | ||
| {ReduceOpType::SumSquare, "output_value_t(0)"}, | ||
| {ReduceOpType::LogSumExp, "output_value_t(-3.4028234663852886e+38f)"}, // log(0) ≈ -inf |
There was a problem hiding this comment.
| {ReduceOpType::LogSumExp, "output_value_t(-3.4028234663852886e+38f)"}, // log(0) ≈ -inf | |
| {ReduceOpType::LogSumExp, "output_value_t(-3.4028234663852886e+38f)"}, // log(0) ≈ -inf |
| {ReduceOpType::LogSumExp, "output_value_t(-3.4028234663852886e+38f)"}, // log(0) ≈ -inf | ||
| {ReduceOpType::L1, "output_value_t(0)"}, | ||
| {ReduceOpType::L2, "output_value_t(0)"}, | ||
| {ReduceOpType::LogSum, "output_value_t(-3.4028234663852886e+38f)"}, // log(0) ≈ -inf |
There was a problem hiding this comment.
| {ReduceOpType::LogSum, "output_value_t(-3.4028234663852886e+38f)"}, // log(0) ≈ -inf | |
| {ReduceOpType::LogSum, "output_value_t(-3.4028234663852886e+38f)"}, // log(0) ≈ -inf |
…inity Replace float literal approximations (-3.4028e+38f) with proper infinity via division-by-zero: output_value_t(-1.0) / output_value_t(0.0) for -inf, output_value_t(1.0) / output_value_t(0.0) for +inf. The float literal approach produced -FLT_MAX instead of -inf, causing test failures where the expected output is -inf (ReduceMax/Min/LogSum empty tensor tests). Division by zero is well-defined in IEEE 754 and WGSL (produces ±inf). Using output_value_t() ensures the expression works for both f32 and f16 output types. Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Some WGSL shader validators reject constant division by zero (output_value_t(-1.0) / output_value_t(0.0)), causing 'Invalid ShaderModule' errors on Windows CI. Replace with bitcast from IEEE 754 bit patterns: bitcast<f32>(0xFF800000u) = -inf bitcast<f32>(0x7F800000u) = +inf bitcast<f32>() is always valid WGSL (u32→f32 reinterpretation). output_value_t() then converts f32 inf to f16 inf when the output type is half-precision (f16 has an infinity representation). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
|
Dmitri Smirnov (@yuslepukhin) Guenther Schmuelling (@guschmue) Do you have suggestions on a fix for webgpu kernel? I tried a few things to produce inf but they don't seem to be liked by the webgpu runtime. Thanks! |
Description
Remove the overly strict assertion in CUDA
PrepareForReducethat rejects reducing along a zero-sized dimension even with explicit axes. This matches the behavior of the CPU implementation which handles empty tensors viacheck_and_reduce_empty_set_input().Motivation
ORT GenAI's Gemma4 CUDA pipeline triggers ReduceSum on
{1, 0}tensors during prefill (past_sequence_length=0). The CPU implementation handles this correctly, but the CUDA path crashes with:Reducing axis 1 of
{1, 0}with keepdims=false produces shape{1}filled with the identity value (0 for sum). This is mathematically valid and numpy handles it correctly.Changes
Removed the
ORT_ENFORCE(input_dims[axis] != 0, ...)assertion at line 291 ofreduction_ops.cc. The existingReduceComputeCorealready handlesinput_count == 0correctly (line 369-370).The default-axes path (line 302) is left unchanged — it already conditionally checks
keepdims || dim != 0.Testing
Verified with Gemma4 e2b-it model on H200 GPU:
ReduceSum_node_232crashes on{1, 0}tensorRelated issues