Additional validation in LSTM and DynamicQuantizedLSTM

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<Tensor>(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<Tensor>(8);
   const Tensor* w_zp = context->Input<Tensor>(9);
   const Tensor* r_scale = context->Input<Tensor>(10);

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<Tensor>(3);              // bias. [num_directions, 8*hidden_size]
+    const auto* sequence_lens = context->Input<Tensor>(4);  // [batch_size]
+    const auto* initial_h = context->Input<Tensor>(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<float>() : nullptr;
     const auto* recurrent_weights = (R != nullptr) ? R->Data<float>() : nullptr;
 

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();

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

onnxruntime/core/providers/cuda/rnn/cudnn_rnn_base.cc

@@ -110,6 +110,19 @@ Status CudnnRnnBase<T>::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<int64_t>(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<int64_t>(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<int64_t>(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<CudaT>()));

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<uint32_t>(hidden_size_);
 
   // Output shapes

onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc

@@ -362,6 +362,71 @@ TEST(DynamicQuantLSTMTest, LargeSize) {
   RunQuantLSTM<uint8_t>(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<std::vector<std::string>>("activations", {"sigmoid", "tanh", "tanh"});
+  test.AddAttribute("direction", "forward");
+  test.AddAttribute("hidden_size", hidden_size);
+  test.AddAttribute<int64_t>("input_forget", 0);
+
+  // X
+  std::vector<int64_t> X_dims = {seq_len, batch_size, input_size};
+  std::vector<float> X_data(seq_len * batch_size * input_size, 0.1f);
+  test.AddInput<float>("X", X_dims, X_data);
+
+  // W: [num_directions, input_size, 4*actual_hidden] - inconsistent with the hidden_size attribute.
+  std::vector<int64_t> W_dims = {num_directions, input_size, 4 * actual_hidden};
+  std::vector<uint8_t> w_quant(num_directions * input_size * 4 * actual_hidden, 1);
+  test.AddInput<uint8_t>("W", W_dims, w_quant);
+
+  // R: [num_directions, actual_hidden, 4*actual_hidden]
+  std::vector<int64_t> R_dims = {num_directions, actual_hidden, 4 * actual_hidden};
+  std::vector<uint8_t> r_quant(num_directions * actual_hidden * 4 * actual_hidden, 1);
+  test.AddInput<uint8_t>("R", R_dims, r_quant);
+
+  // B, sequence_lens
+  test.AddOptionalInputEdge<float>();
+  test.AddOptionalInputEdge<int>();
+
+  // initial_h / initial_c (sized using the hidden_size attribute)
+  std::vector<int64_t> initial_dims = {num_directions, batch_size, hidden_size};
+  std::vector<float> initial_data(num_directions * batch_size * hidden_size, 0.0f);
+  test.AddInput<float>("initial_h", initial_dims, initial_data);
+  test.AddInput<float>("initial_c", initial_dims, initial_data);
+
+  // P
+  test.AddOptionalInputEdge<float>();
+
+  // Per-tensor quantization parameters (their shapes do not involve hidden_size).
+  std::vector<int64_t> per_tensor_dims = {num_directions};
+  test.AddInput<float>("W_scale", per_tensor_dims, std::vector<float>(num_directions, 1.0f));
+  test.AddInput<uint8_t>("W_zero_point", per_tensor_dims, std::vector<uint8_t>(num_directions, 0));
+  test.AddInput<float>("R_scale", per_tensor_dims, std::vector<float>(num_directions, 1.0f));
+  test.AddInput<uint8_t>("R_zero_point", per_tensor_dims, std::vector<uint8_t>(num_directions, 0));
+
+  // Outputs (sized using the hidden_size attribute); values are not checked because we expect failure.
+  std::vector<int64_t> Y_dims = {seq_len, num_directions, batch_size, hidden_size};
+  test.AddOutput<float>("Y", Y_dims, std::vector<float>(seq_len * num_directions * batch_size * hidden_size, 0.0f));
+  std::vector<int64_t> Y_h_dims = {num_directions, batch_size, hidden_size};
+  test.AddOutput<float>("Y_h", Y_h_dims, std::vector<float>(num_directions * batch_size * hidden_size, 0.0f));
+  std::vector<int64_t> Y_c_dims = {num_directions, batch_size, hidden_size};
+  test.AddOutput<float>("Y_c", Y_c_dims, std::vector<float>(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) {

onnxruntime/test/providers/cpu/rnn/deep_cpu_lstm_op_test.cc

@@ -1349,6 +1349,55 @@ 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<std::vector<std::string>>("activations", {"sigmoid", "tanh", "tanh"});
+  test.AddAttribute("direction", "forward");
+  test.AddAttribute("hidden_size", hidden_size);
+
+  std::vector<int64_t> X_dims = {seq_length, batch_size, input_size};
+  std::vector<float> X_data(seq_length * batch_size * input_size, 0.1f);
+  test.AddInput<float>("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<int64_t> W_dims = {1, 4 * bogus_w_hidden, input_size};
+  std::vector<float> W_data(1 * 4 * bogus_w_hidden * input_size, 0.1f);
+  test.AddInput<float>("W", W_dims, W_data);
+
+  std::vector<int64_t> R_dims = {1, 4 * hidden_size, hidden_size};
+  std::vector<float> R_data(1 * 4 * hidden_size * hidden_size, 0.1f);
+  test.AddInput<float>("R", R_dims, R_data);
+
+  // Optional inputs left empty.
+  test.AddOptionalInputEdge<float>();  // B
+  test.AddOptionalInputEdge<int>();    // sequence_lens
+  test.AddOptionalInputEdge<float>();  // initial_h
+  test.AddOptionalInputEdge<float>();  // initial_c
+  test.AddOptionalInputEdge<float>();  // P
+
+  std::vector<int64_t> Y_dims = {seq_length, 1, batch_size, hidden_size};
+  std::vector<float> Y_data(seq_length * 1 * batch_size * hidden_size, 0.f);
+  test.AddOutput<float>("Y", Y_dims, Y_data);
+
+  // 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, kDmlExecutionProvider, kQnnExecutionProvider});
+}
+
 #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) {