Updates to support HF dynamic cache export

backends/mlx/builder/op_helpers.py

@@ -187,6 +187,7 @@ def emit_shape(
     slot: Slot,
     *,
     end_dim: "Optional[int]" = None,
+    dim_offset: int = 0,
 ) -> "list[IntOrVid]":
     """Return the shape of ``node`` as a list of ``IntOrVid``.
 
@@ -197,11 +198,13 @@ def emit_shape(
     Args:
         P: program builder.
         node: FX node whose shape to walk (must have ``meta['val']``).
-        slot: slot corresponding to ``node`` (used as the
-            ``SymSize`` source for any dynamic dim).
+        slot: tensor slot used as the ``SymSize`` source for dynamic dims.
         end_dim: stop index (exclusive). ``None`` means the full ndim.
             Negative values index from the end (e.g. ``-1`` is "all
             leading dims, drop the last").
+        dim_offset: offset added to dynamic dimension indices when reading
+            ``slot``. Use when its runtime axes differ from the metadata axes;
+            static dimensions still use the metadata values.
 
     Returns:
         ``list[IntOrVid]`` of length ``end_dim`` (after normalization).
@@ -228,7 +231,7 @@ def emit_shape(
             P.emit(
                 SymSizeNode(
                     a=P.slot_to_tid(slot),
-                    dim=dim_idx,
+                    dim=dim_idx + dim_offset,
                     out=P.slot_to_vid(d_val),
                 )
             )

backends/mlx/ops.py

@@ -2466,8 +2466,7 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     indices = [P.slot_to_tid(idx) for _, idx in non_none]
     axes = [i for i, _ in non_none]
 
-    # slice_sizes: 1 for indexed axes, full dim size for non-indexed axes
-    # Use int() to handle SymInt values from dynamic shapes
+    # slice_sizes: 1 for indexed axes, full static size for non-indexed axes.
     indexed_axes = set(axes)
     slice_sizes = []
     for dim in range(x_ndim):
@@ -2495,18 +2494,47 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     )
 
     # Reshape to match aten.index.Tensor output shape, which strips the
-    # trailing dimensions introduced by gather's slice_sizes
+    # trailing dimensions introduced by gather's slice_sizes.
     out_meta = n.meta.get("val")
     if out_meta is None:
         raise ValueError(
             "aten.index.Tensor: output shape metadata required for reshape after gather"
         )
-    out_shape = [P.to_int_or_vid(int(d)) for d in out_meta.shape]
+    # Non-indexed sizes are static above, so symbolic output sizes belong to
+    # the broadcast index shape, which leads the gather result. Read them at
+    # runtime instead of adding specialization guards through int(SymInt).
+    # Contiguous indexed axes keep the broadcast dimensions in place in ATen.
+    leading_dims = axes[0] if axes == list(range(axes[0], axes[-1] + 1)) else 0
+    # Read pre-transpose sizes; the offset maps broadcast dims back to gather axes.
+    out_shape = emit_shape(P, n, gather_slot, dim_offset=-leading_dims)
+
+    reshape_slot = gather_slot
+    broadcast_ndim = len(out_meta.shape) - x_ndim + len(axes)
+    broadcast_shape = out_meta.shape[leading_dims : leading_dims + broadcast_ndim]
+    # Moving singleton blocks does not change element order; keep those
+    # lowerings reshape-only, without extra instructions or tensor slots.
+    if any(size > 1 for size in slice_sizes[:leading_dims]) and any(
+        not isinstance(size, int) or size > 1 for size in broadcast_shape
+    ):
+        _, reshape_slot = P.make_tmp_slot()
+        P.emit(
+            TransposeNode(
+                x=P.slot_to_tid(gather_slot),
+                out=P.slot_to_tid(reshape_slot),
+                perm=(
+                    list(range(broadcast_ndim, broadcast_ndim + leading_dims))
+                    + list(range(broadcast_ndim))
+                    + list(
+                        range(broadcast_ndim + leading_dims, broadcast_ndim + x_ndim)
+                    )
+                ),
+            )
+        )
 
     out = P.make_or_get_slot(n)
     P.emit(
         ReshapeNode(
-            x=P.slot_to_tid(gather_slot),
+            x=P.slot_to_tid(reshape_slot),
             out=P.slot_to_tid(out),
             shape=out_shape,
         )

backends/mlx/passes.py

@@ -216,8 +216,9 @@ class CollapseDtypeConversionPass(ExportPass):
 
     _to_copy(dtype=bf16)(_to_copy(dtype=f32)(x)) → _to_copy(dtype=bf16)(x)
 
-    Only the final dtype matters. Only collapses when both nodes are pure dtype
-    conversions (no device/layout/memory_format changes).
+    Only collapse when the intermediate cast preserves every source value.
+    Narrowing or cross-kind casts may round, truncate, or overflow and must stay.
+    Both nodes must be pure dtype conversions (no device/layout/memory_format changes).
     """
 
     def call(self, graph_module: GraphModule) -> PassResult:
@@ -244,8 +245,31 @@ def call(self, graph_module: GraphModule) -> PassResult:
             if not _is_pure_dtype_cast(node_kw) or not _is_pure_dtype_cast(parent_kw):
                 continue
 
+            source = parent.args[0]
+            source_val = source.meta.get("val") if isinstance(source, Node) else None
+            if source_val is None:
+                continue
+            source_dtype = source_val.dtype
+            intermediate_dtype = parent_kw["dtype"]
+            if source_dtype != intermediate_dtype and (
+                source_dtype,
+                intermediate_dtype,
+            ) not in {
+                # Boolean values 0 and 1 are exact in each floating-point dtype.
+                (torch.bool, torch.float16),
+                (torch.bool, torch.bfloat16),
+                (torch.bool, torch.float32),
+                (torch.bool, torch.float64),
+                (torch.float16, torch.float32),
+                (torch.bfloat16, torch.float32),
+                (torch.float16, torch.float64),
+                (torch.bfloat16, torch.float64),
+                (torch.float32, torch.float64),
+            }:
+                continue
+
             # Rewrite: to_copy(to_copy(x, dtype=d1), dtype=d2) → to_copy(x, dtype=d2)
-            node.args = (parent.args[0],)
+            node.args = (source,)
             graph.erase_node(parent)
             modified = True
 

backends/mlx/runtime/MLXInterpreter.h

@@ -12,6 +12,7 @@
 #include "MLXExecutor.h"
 
 #include <algorithm>
+#include <limits>
 #include <vector>
 
 #include <mlx/array.h>
@@ -307,6 +308,17 @@ inline void exec_sdpa(const SdpaNode& n, ExecutionState& st, StreamOrDevice s) {
       sinks,
       false,
       s);
+  if (n.mask) {
+    const auto& M = st.const_tensor_ref(*n.mask);
+    array allowed = M.dtype() == bool_
+        ? M
+        : not_equal(
+              M, array(-std::numeric_limits<float>::infinity(), M.dtype()), s);
+    array has_key = M.ndim() == 0 ? allowed : any(allowed, -1, true, s);
+    // MLX can return nonzero values or NaNs for empty rows. Select, rather
+    // than multiply, to preserve PyTorch's zero-row semantics in both cases.
+    out = where(has_key, out, array(0, out.dtype()), s);
+  }
   st.set_tensor(n.out, std::move(out));
 }