@@ -107,14 +107,56 @@ struct FuseAttention final : public PredicateBasedPass {
     if (!CheckKind(dim, kSlice)) {
       return false;
     }
-    // The slice must select the last dimension of the shape.
+    // `q` is rank 4 (checked by the caller), so the slice must select exactly
+    // the last dimension: starts=[-1], ends >= 4, optional axes in {0, -1} and
+    // optional steps=1. Anything else could select more than one dimension.
+    const size_t num_inputs = dim->inputs().size();
     int64_t start;
-    if (dim->inputs().size() < 2 ||
-        !FetchSoleIntValueOfTensor(dim->input(1), start) || start != -1) {
+    int64_t end;
+    if (num_inputs < 3 || num_inputs > 5 ||
+        !FetchSoleIntValueOfTensor(dim->input(1), start) || start != -1 ||
+        !FetchSoleIntValueOfTensor(dim->input(2), end) || end < 4) {
       return false;
     }
+    if (num_inputs >= 4) {
+      int64_t axis;
+      if (!FetchSoleIntValueOfTensor(dim->input(3), axis) ||
+          (axis != 0 && axis != -1)) {
+        return false;
+      }
+    }
+    if (num_inputs == 5) {
+      int64_t step;
+      if (!FetchSoleIntValueOfTensor(dim->input(4), step) || step != 1) {
+        return false;
+      }
+    }
     const Node* shape = dim->input(0)->node();
-    return CheckKind(shape, "Shape") && shape->input(0) == q;
+    if (!CheckKind(shape, "Shape") || shape->input(0) != q) {
+      return false;
+    }
+    // The Shape must cover all of `q`; a start/end selection would shift which
+    // dimension the slice reads.
+    const auto range = FetchStartAndEndAttrOfShape(shape, 4);
+    return range.first == 0 && range.second == 4;
+  }
+
+  // Returns true when `v` is a rank-4 `[batch, heads, seq, head_size]` tensor
+  // whose batch and head dimensions are statically known.
+  static bool IsRank4AttentionLayout(const Value* v) {
+    if (!v->has_sizes() || v->sizes().size() != 4) {
+      return false;
+    }
+    const auto& dims = v->sizes();
+    return dims[0].is_int && dims[1].is_int;
+  }
+
+  // Returns true when `a` and `b` agree on their (static) batch and head
+  // dimensions. Attention requires matching batch and head counts for Q, K, V.
+  static bool HaveSameBatchAndHeads(const Value* a, const Value* b) {
+    const auto& da = a->sizes();
+    const auto& db = b->sizes();
+    return da[0].dim == db[0].dim && da[1].dim == db[1].dim;
   }
 
   bool runTransform(Node* n, Graph& graph,
@@ -215,6 +257,14 @@ struct FuseAttention final : public PredicateBasedPass {
       return false;
     }
 
+    // Attention only represents rank-4 [batch, heads, seq, head_size] inputs
+    // with matching batch and head counts. Leave every other graph unchanged.
+    if (!IsRank4AttentionLayout(query) || !IsRank4AttentionLayout(key) ||
+        !IsRank4AttentionLayout(value) || !HaveSameBatchAndHeads(query, key) ||
+        !HaveSameBatchAndHeads(query, value)) {
+      return false;
+    }
+
     // Resolve the effective scale applied to Q @ K^T.
     bool has_scale = false;
     float scale = 0.0f;
@@ -223,6 +273,14 @@ struct FuseAttention final : public PredicateBasedPass {
     if (FetchSoleValueOfTensor(scale_q, scale_q_val) &&
         FetchSoleValueOfTensor(scale_k, scale_k_val)) {
       // Explicit scale: torch emits sqrt(scale) constants on both branches.
+      // Only scalars are accepted: a higher-rank scale would broadcast Q or K^T
+      // to a larger rank than the fused Attention output.
+      auto is_scalar = [](const Value* v) {
+        return v->has_sizes() && v->sizes().empty();
+      };
+      if (!is_scalar(scale_q) || !is_scalar(scale_k)) {
+        return false;
+      }
       has_scale = true;
       scale = scale_q_val * scale_k_val;
     } else if (CheckKind(scale_q, kSqrt) && CheckKind(scale_k, kSqrt) &&
 
@@ -5105,6 +5105,37 @@ def vi(name, shape):
         optimized_model = self._optimized(model, ["fuse_attention"])
         assert all(n.op_type != "Attention" for n in optimized_model.graph.node)
 
+    def test_fuse_attention_shape_end_selection_no_fuse(self):  # type: () -> None
+        # Shape(q, end=3) followed by Slice(-1) reads dim 2, not head_size, so
+        # the default scale must not be assumed.
+        model = self._make_sdpa_model(mask="none", scale=None)
+        shape_node = [n for n in model.graph.node if n.op_type == "Shape"][0]
+        shape_node.attribute.append(helper.make_attribute("end", 3))
+        optimized_model = self._optimized(model, ["fuse_attention"])
+        assert all(n.op_type != "Attention" for n in optimized_model.graph.node)
+
+    def test_fuse_attention_leading_dim_scale_no_fuse(self):  # type: () -> None
+        # A [1,1,1,1,1] scale would promote Q/K^T to rank 5; the fused Attention
+        # output would then have the wrong rank.
+        model = self._make_sdpa_model(mask="none", scale=0.125)
+        sf = float(np.sqrt(0.125))
+        for i, init in enumerate(model.graph.initializer):
+            if init.name in ("sqf_q", "sqf_k"):
+                model.graph.initializer[i].CopyFrom(
+                    numpy_helper.from_array(
+                        np.full((1, 1, 1, 1, 1), sf, np.float32), init.name
+                    )
+                )
+        optimized_model = self._optimized(model, ["fuse_attention"])
+        assert all(n.op_type != "Attention" for n in optimized_model.graph.node)
+
+    def test_fuse_attention_batch_mismatch_no_fuse(self):  # type: () -> None
+        # Q has batch 1 while K/V have batch 2; Attention needs matching batches.
+        model = self._make_sdpa_model(mask="none", scale=None)
+        model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1
+        optimized_model = self._optimized(model, ["fuse_attention"])
+        assert all(n.op_type != "Attention" for n in optimized_model.graph.node)
+
     @unittest.skipUnless(has_torch, "onnx test needs torch")
     def test_fuse_attention_torch_sdpa_export(self):  # type: () -> None
         # Fuse graphs produced by the real torch.onnx exporter.