fix mha test

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/_smollm_2.py

@@ -459,7 +459,7 @@ def get_ort_inputs(self):
         if not hasattr(self, "_ort_inputs"):
             inputs = {
                 "input_ids": numpy.random.randint(0, 49152, (1, 30)).astype(numpy.int64),
-                "position_ids": numpy.ones((1, 30), dtype=numpy.int64),
+                "position_ids": numpy.arange(30).reshape(1, 30).astype(numpy.int64),
                 "past_key_values_0_0": numpy.random.rand(1, 32, 16, 64).astype(numpy.float32),
                 "past_key_values_0_1": numpy.random.rand(1, 32, 16, 64).astype(numpy.float32),
             }

onnxscript/rewriter/ort_fusions/cos_sin_cache.py

@@ -97,7 +97,7 @@ def _compute_dynamic_freqs(self, op, inv_freq, position_ids, dtype):
         # If we do not compute max_pos_id using ONNX ops, use inv_freq and position_ids
         # to compute angles and cos/sin values
         inv_freq_values = inv_freq.const_value.numpy().reshape(1, -1)
-        pos_id_range = np.arange(max_pos_id, dtype=np.float32).reshape(-1, 1)
+        pos_id_range = np.arange(max_pos_id + 1, dtype=np.float32).reshape(-1, 1)
         angles = np.matmul(pos_id_range, inv_freq_values)
         return self._compute_const_freqs(op, angles)