From a8f0afaa7a8167563fdbbbf0547ca1941c3648e4 Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Tue, 7 Jul 2026 18:46:14 +1000 Subject: [PATCH 1/2] Validate weights against hidden_size attribute for LSTM and DynamicQuantizedLSTM --- .../cpu/quantization/dynamic_quantize_lstm.cc | 15 +++++ .../core/providers/cpu/rnn/deep_cpu_lstm.cc | 10 +++ .../core/providers/cpu/rnn/rnn_helpers.cc | 24 ++++--- .../core/providers/cpu/rnn/rnn_helpers.h | 7 +- .../test/contrib_ops/quantize_lstm_op_test.cc | 65 +++++++++++++++++++ .../cpu/rnn/deep_cpu_lstm_op_test.cc | 48 ++++++++++++++ 6 files changed, 159 insertions(+), 10 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc index ba3c8a2ada30b..63ac2d20bedcd 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc @@ -172,6 +172,21 @@ Status DynamicQuantizeLSTM::Compute(OpKernelContext* context) const { const auto& W_shape = (W != nullptr) ? W->Shape() : packed_W_.shape_; const auto& R_shape = (R != nullptr) ? R->Shape() : packed_R_.shape_; + // The WeightCheck macro below only validates the scale/zero-point shapes against hidden_size_, and + // only for per-channel (2D) quantization. It does not validate the W and R weight tensor shapes. + // Validate those here so a model whose hidden_size attribute is inconsistent with the weights is + // rejected cleanly rather than causing an out-of-bounds access in the compute below (ONNX shape + // inference does not verify this relationship). The quantized weight layout is transposed relative + // to the standard LSTM: W is [num_directions, input_size, 4*hidden_size] and + // R is [num_directions, hidden_size, 4*hidden_size], hence weights_transposed=true. The remaining + // inputs (B, sequence_lens, initial_h, initial_c, P) are validated by LSTMBase::ValidateInputs in + // ComputeImpl, so they are passed as nullptr here. + const Tensor& X = *context->Input(0); // inputs. [seq_length, batch_size, input_size] + ORT_RETURN_IF_ERROR(rnn::detail::ValidateCommonRnnInputs(X, W_shape, R_shape, nullptr /*B*/, 4, + nullptr /*sequence_lens*/, nullptr /*initial_h*/, + num_directions_, hidden_size_, + true /*weights_transposed*/)); + const Tensor* w_scale = context->Input(8); const Tensor* w_zp = context->Input(9); const Tensor* r_scale = context->Input(10); diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc index d2520804bb64c..56c1c22b14851 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc @@ -291,6 +291,16 @@ Status DeepCpuLstmOp::Compute(OpKernelContext* context) const { const auto& W_shape = (W != nullptr) ? W->Shape() : packed_W_.shape_; const auto& R_shape = (R != nullptr) ? R->Shape() : packed_R_.shape_; + // Validate that the W and R weight shapes are consistent with the hidden_size attribute and the + // input shape, matching the RNN and GRU kernels. ONNX shape inference does not verify this + // relationship, so an inconsistent model (e.g. a bogus hidden_size) would otherwise reach the + // compute below and cause an out-of-bounds access instead of a clean error. + const auto* B = context->Input(3); // bias. [num_directions, 8*hidden_size] + const auto* sequence_lens = context->Input(4); // [batch_size] + const auto* initial_h = context->Input(5); // initial hidden. [num_directions, batch_size, hidden_size] + ORT_RETURN_IF_ERROR(rnn::detail::ValidateCommonRnnInputs(X, W_shape, R_shape, B, 4, sequence_lens, + initial_h, num_directions_, hidden_size_)); + const auto* input_weights = (W != nullptr) ? W->Data() : nullptr; const auto* recurrent_weights = (R != nullptr) ? R->Data() : nullptr; diff --git a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.cc b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.cc index f522b801fb4c8..6b99b7fd7e02b 100644 --- a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.cc +++ b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.cc @@ -35,7 +35,8 @@ Status ValidateCommonRnnInputs(const Tensor& X, const Tensor* sequence_lens, const Tensor* initial_h, int64_t num_directions, - int64_t hidden_size) { + int64_t hidden_size, + bool weights_transposed) { auto& X_shape = X.Shape(); int64_t seq_length = X_shape[0]; @@ -45,21 +46,26 @@ Status ValidateCommonRnnInputs(const Tensor& X, if (X_shape.NumDimensions() != 3) return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input X must have 3 dimensions only. Actual:", X_shape); + // The gate dimension (mult * hidden_size) and the input/hidden dimension are on dims 1 and 2 + // respectively for the standard layout, and swapped for the transposed (quantized) layout. + const int64_t gate_dim = hidden_size * WRB_dim_1_multipler; if (W_shape.NumDimensions() != 3 || W_shape[0] != num_directions || - W_shape[1] != hidden_size * WRB_dim_1_multipler || - W_shape[2] != input_size) + (weights_transposed ? W_shape[2] : W_shape[1]) != gate_dim || + (weights_transposed ? W_shape[1] : W_shape[2]) != input_size) return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input W must have shape {", - num_directions, ",", WRB_dim_1_multipler, "*", hidden_size, ",", - input_size, "}. Actual:", W_shape); + num_directions, ",", + weights_transposed ? input_size : gate_dim, ",", + weights_transposed ? gate_dim : input_size, "}. Actual:", W_shape); if (R_shape.NumDimensions() != 3 || R_shape[0] != num_directions || - R_shape[1] != hidden_size * WRB_dim_1_multipler || - R_shape[2] != hidden_size) + (weights_transposed ? R_shape[2] : R_shape[1]) != gate_dim || + (weights_transposed ? R_shape[1] : R_shape[2]) != hidden_size) return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input R must have shape {", - num_directions, ",", WRB_dim_1_multipler, "*", hidden_size, ",", - hidden_size, "}. Actual:", R_shape); + num_directions, ",", + weights_transposed ? hidden_size : gate_dim, ",", + weights_transposed ? gate_dim : hidden_size, "}. Actual:", R_shape); if (B != nullptr) { auto& B_shape = B->Shape(); diff --git a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h index 1abd9b7408f12..188bea7c69551 100644 --- a/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h +++ b/onnxruntime/core/providers/cpu/rnn/rnn_helpers.h @@ -100,7 +100,12 @@ Status ValidateCommonRnnInputs(const Tensor& X, const Tensor* sequence_lens, const Tensor* initial_h, int64_t num_directions, - int64_t hidden_size); + int64_t hidden_size, + // When true, the W and R gate dimension and input/hidden dimension are + // swapped, i.e. W is [num_directions, input_size, mult*hidden_size] and + // R is [num_directions, hidden_size, mult*hidden_size]. This is the layout + // used by the quantized LSTM (DynamicQuantizeLSTM). + bool weights_transposed = false); /// Copy an input array repeatedly to an output array /// @param input_begin Beginning of input diff --git a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc index 93a63ce19eb03..704c24b2f12e3 100644 --- a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc +++ b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc @@ -362,6 +362,71 @@ TEST(DynamicQuantLSTMTest, LargeSize) { RunQuantLSTM(12, 3, 278); } +// Regression test: a DynamicQuantizeLSTM whose hidden_size attribute is inconsistent with the actual +// W/R weight tensor shapes must be rejected cleanly. With per-tensor quantization the scale/zero-point +// shape checks do not involve hidden_size, and ONNX shape inference does not verify the weight shapes, +// so such a model would otherwise reach the compute path and cause an out-of-bounds access. +TEST(DynamicQuantLSTMTest, MismatchedWeightShapeIsRejected) { + OpTester test("DynamicQuantizeLSTM", 1 /*opset_version*/, onnxruntime::kMSDomain /*domain*/); + + constexpr int num_directions = 1; + constexpr int64_t seq_len = 1; + constexpr int64_t batch_size = 1; + constexpr int64_t input_size = 2; + constexpr int64_t actual_hidden = 2; // the weights are sized for this hidden size + constexpr int64_t hidden_size = 3; // bogus attribute value, inconsistent with the weights + + test.AddAttribute>("activations", {"sigmoid", "tanh", "tanh"}); + test.AddAttribute("direction", "forward"); + test.AddAttribute("hidden_size", hidden_size); + test.AddAttribute("input_forget", 0); + + // X + std::vector X_dims = {seq_len, batch_size, input_size}; + std::vector X_data(seq_len * batch_size * input_size, 0.1f); + test.AddInput("X", X_dims, X_data); + + // W: [num_directions, input_size, 4*actual_hidden] - inconsistent with the hidden_size attribute. + std::vector W_dims = {num_directions, input_size, 4 * actual_hidden}; + std::vector w_quant(num_directions * input_size * 4 * actual_hidden, 1); + test.AddInput("W", W_dims, w_quant); + + // R: [num_directions, actual_hidden, 4*actual_hidden] + std::vector R_dims = {num_directions, actual_hidden, 4 * actual_hidden}; + std::vector r_quant(num_directions * actual_hidden * 4 * actual_hidden, 1); + test.AddInput("R", R_dims, r_quant); + + // B, sequence_lens + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + + // initial_h / initial_c (sized using the hidden_size attribute) + std::vector initial_dims = {num_directions, batch_size, hidden_size}; + std::vector initial_data(num_directions * batch_size * hidden_size, 0.0f); + test.AddInput("initial_h", initial_dims, initial_data); + test.AddInput("initial_c", initial_dims, initial_data); + + // P + test.AddOptionalInputEdge(); + + // Per-tensor quantization parameters (their shapes do not involve hidden_size). + std::vector per_tensor_dims = {num_directions}; + test.AddInput("W_scale", per_tensor_dims, std::vector(num_directions, 1.0f)); + test.AddInput("W_zero_point", per_tensor_dims, std::vector(num_directions, 0)); + test.AddInput("R_scale", per_tensor_dims, std::vector(num_directions, 1.0f)); + test.AddInput("R_zero_point", per_tensor_dims, std::vector(num_directions, 0)); + + // Outputs (sized using the hidden_size attribute); values are not checked because we expect failure. + std::vector Y_dims = {seq_len, num_directions, batch_size, hidden_size}; + test.AddOutput("Y", Y_dims, std::vector(seq_len * num_directions * batch_size * hidden_size, 0.0f)); + std::vector Y_h_dims = {num_directions, batch_size, hidden_size}; + test.AddOutput("Y_h", Y_h_dims, std::vector(num_directions * batch_size * hidden_size, 0.0f)); + std::vector Y_c_dims = {num_directions, batch_size, hidden_size}; + test.AddOutput("Y_c", Y_c_dims, std::vector(num_directions * batch_size * hidden_size, 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, "Input W must have shape"); +} + #ifndef ENABLE_TRAINING // Prepacking is disabled in full training build so no need to test the feature in a training build. TEST(DynamicQuantLSTMTest, SharedPrepackedWeights) { diff --git a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc index c23e843b6f143..af5753c6e5e84 100644 --- a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc +++ b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc @@ -1349,6 +1349,54 @@ TEST(LSTMTest, ONNXRuntime_TestLSTMZeroSeqInMiddle) { &sequence_length, use_bias, use_peepholes, 0.0f, false, false); } +// Regression test for a model whose hidden_size attribute is inconsistent with the actual W +// weight tensor shape. ONNX shape inference/checker does not validate the relationship between +// hidden_size and the W/R weight shapes, so such a model previously reached the LSTM kernel and +// caused an out-of-bounds access. The kernel must now reject it up front via ValidateCommonRnnInputs, +// matching the behavior of the RNN and GRU kernels. +TEST(LSTMTest, MismatchedWeightShapeIsRejected) { + OpTester test("LSTM"); + + constexpr int64_t seq_length = 1; + constexpr int64_t batch_size = 1; + constexpr int64_t input_size = 2; + constexpr int64_t hidden_size = 2; // attribute value; sizes the outputs and R + constexpr int64_t bogus_w_hidden = 3; // W is sized for a different (larger) hidden size + + test.AddAttribute>("activations", {"sigmoid", "tanh", "tanh"}); + test.AddAttribute("direction", "forward"); + test.AddAttribute("hidden_size", hidden_size); + + std::vector X_dims = {seq_length, batch_size, input_size}; + std::vector X_data(seq_length * batch_size * input_size, 0.1f); + test.AddInput("X", X_dims, X_data); + + // W has 4 * bogus_w_hidden rows instead of 4 * hidden_size, which is inconsistent with the + // hidden_size attribute. + std::vector W_dims = {1, 4 * bogus_w_hidden, input_size}; + std::vector W_data(1 * 4 * bogus_w_hidden * input_size, 0.1f); + test.AddInput("W", W_dims, W_data); + + std::vector R_dims = {1, 4 * hidden_size, hidden_size}; + std::vector R_data(1 * 4 * hidden_size * hidden_size, 0.1f); + test.AddInput("R", R_dims, R_data); + + // Optional inputs left empty. + test.AddOptionalInputEdge(); // B + test.AddOptionalInputEdge(); // sequence_lens + test.AddOptionalInputEdge(); // initial_h + test.AddOptionalInputEdge(); // initial_c + test.AddOptionalInputEdge(); // P + + std::vector Y_dims = {seq_length, 1, batch_size, hidden_size}; + std::vector Y_data(seq_length * 1 * batch_size * hidden_size, 0.f); + test.AddOutput("Y", Y_dims, Y_data); + + // Expect a clean validation error rather than a crash. + test.Run(OpTester::ExpectResult::kExpectFailure, "Input W must have shape", + {kTensorrtExecutionProvider}); +} + #ifndef ENABLE_TRAINING // Prepacking is disabled in full training build so no need to test the feature in a training build. TEST(LSTMTest, ONNXRuntime_TestLSTMForward_OpSet22_CUDA) { From b19f5c5bfab8e42018741887dacdd1a3d2cab9a0 Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Wed, 8 Jul 2026 19:45:49 +1000 Subject: [PATCH 2/2] Add check to WebGPU and CUDA implementations Skip DML and QNN as they're no longer maintained in this repo. --- .../core/providers/cuda/rnn/cudnn_rnn_base.cc | 13 +++++++++++++ onnxruntime/core/providers/webgpu/rnn/lstm.cc | 14 ++++++++++++++ .../providers/cpu/rnn/deep_cpu_lstm_op_test.cc | 5 +++-- 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc index 62ec59c0ef798..7cec77c617fb8 100644 --- a/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc +++ b/onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc @@ -110,6 +110,19 @@ Status CudnnRnnBase::ReorganizeWeights(const Tensor* W, const Tensor* R, cons // LSTM R[num_directions_, 4*hidden_size_, hidden_size_] // LSTM B[num_directions_, 8*hidden_size_] size_t number = W_lin_layer_id_.size(); + // Validate that the W and R gate dimension (number * hidden_size_) and R's hidden dimension are + // consistent with the hidden_size attribute. ONNX shape inference does not verify this relationship, + // so an inconsistent model (e.g. a bogus hidden_size) would otherwise be used to size the reorganized + // weight buffer below and cause an out-of-bounds read from the W/R data. + ORT_RETURN_IF(W->Shape()[0] != num_directions_ || + W->Shape()[1] != SafeInt(number) * hidden_size_, + "Input W must have shape {", num_directions_, ", ", number, "*", hidden_size_, + ", input_size}. Actual: ", W->Shape()); + ORT_RETURN_IF(R->Shape()[0] != num_directions_ || + R->Shape()[1] != SafeInt(number) * hidden_size_ || + R->Shape()[2] != hidden_size_, + "Input R must have shape {", num_directions_, ", ", number, "*", hidden_size_, + ", ", hidden_size_, "}. Actual: ", R->Shape()); int64_t w_size = SafeInt(num_directions_) * number * hidden_size_ * (input_size + hidden_size_ + 2); TensorShapeVector dims_w({w_size, 1, 1}); ORT_RETURN_IF_ERROR(target_w_desc.Set(dims_w, CudnnTensor::GetDataType())); diff --git a/onnxruntime/core/providers/webgpu/rnn/lstm.cc b/onnxruntime/core/providers/webgpu/rnn/lstm.cc index 1c30ab4300f44..871e96b671d82 100644 --- a/onnxruntime/core/providers/webgpu/rnn/lstm.cc +++ b/onnxruntime/core/providers/webgpu/rnn/lstm.cc @@ -317,6 +317,20 @@ Status Lstm::ComputeInternal(ComputeContext& context) const { input_size = X_shape[2]; } + // Validate that the W and R weight shapes are consistent with the hidden_size attribute. ONNX shape + // inference does not verify this relationship, so an inconsistent model (e.g. a bogus hidden_size) + // would otherwise be used to drive the cell computation and read out of bounds from the W/R buffers. + const auto& W_shape = W->Shape(); + const auto& R_shape = R->Shape(); + ORT_RETURN_IF(W_shape.NumDimensions() != 3 || W_shape[0] != num_directions || + W_shape[1] != 4 * hidden_size_ || W_shape[2] != input_size, + "LSTM: Input W must have shape {", num_directions, ", 4*", hidden_size_, ", ", input_size, + "}. Actual: ", W_shape); + ORT_RETURN_IF(R_shape.NumDimensions() != 3 || R_shape[0] != num_directions || + R_shape[1] != 4 * hidden_size_ || R_shape[2] != hidden_size_, + "LSTM: Input R must have shape {", num_directions, ", 4*", hidden_size_, ", ", hidden_size_, + "}. Actual: ", R_shape); + uint32_t H = static_cast(hidden_size_); // Output shapes diff --git a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc index af5753c6e5e84..a2b4e0a78fcd9 100644 --- a/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc +++ b/onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc @@ -1392,9 +1392,10 @@ TEST(LSTMTest, MismatchedWeightShapeIsRejected) { std::vector Y_data(seq_length * 1 * batch_size * hidden_size, 0.f); test.AddOutput("Y", Y_dims, Y_data); - // Expect a clean validation error rather than a crash. + // Expect a clean validation error rather than a crash. Skip DML (deprecated) and QNN (now maintained + // out-of-tree in the onnxruntime-qnn repo), which do not perform this validation. test.Run(OpTester::ExpectResult::kExpectFailure, "Input W must have shape", - {kTensorrtExecutionProvider}); + {kTensorrtExecutionProvider, kDmlExecutionProvider, kQnnExecutionProvider}); } #ifndef ENABLE_TRAINING