Add stricter condition to matrix sdpa

docs/src/python/optimizers.rst

@@ -51,14 +51,14 @@ the saved state. Here's a simple example:
    optimizer.update(model, grads)
 
    # Save the state
-   state = tree_flatten(optimizer.state, destination={})
-   mx.save_safetensors("optimizer.safetensors", state)
+   state = tree_flatten(optimizer.state)
+   mx.save_safetensors("optimizer.safetensors", dict(state))
 
    # Later on, for example when loading from a checkpoint,
    # recreate the optimizer and load the state
    optimizer = optim.Adam(learning_rate=1e-2)
 
-   state = tree_unflatten(mx.load("optimizer.safetensors"))
+   state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
    optimizer.state = state
 
 Note, not every optimizer configuation parameter is saved in the state. For

docs/src/usage/export.rst

@@ -151,7 +151,7 @@ parameters, pass them as inputs to the ``call`` wrapper:
      model.update(tree_unflatten(list(params.items())))
      return model(x)
  
-   params = tree_flatten(model.parameters(), destination={})
+   params = dict(tree_flatten(model.parameters()))
    mx.export_function("model.mlxfn", call, (mx.zeros(4),), params)
 
 

mlx/backend/cuda/scaled_dot_product_attention.cu

@@ -8,13 +8,19 @@
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/dtype_utils.h"
 #include "mlx/fast_primitives.h"
-#include "mlx/transforms_impl.h"
 
+// cudnn_frontend.h redefines this macro.
+#undef CHECK_CUDA_ERROR
+
+#include <cudnn_frontend.h>
+#include <fmt/format.h>
 #include <nvtx3/nvtx3.hpp>
 
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 
+namespace fe = cudnn_frontend;
+
 namespace mlx::core {
 
 namespace cu {
@@ -639,6 +645,294 @@ void sdpa_vector_fallback(
   }
 }
 
+struct SDPACacheKey {
+  int device_id;
+  fe::DataType_t cudnn_type;
+
+  int B;
+  int H;
+  int D;
+
+  int qL;
+  int kL;
+
+  int gqa_factor;
+  float scale;
+
+  int64_t Q_strides[3];
+  int64_t K_strides[3];
+  int64_t V_strides[3];
+  int64_t O_strides[3];
+
+  bool generate_stats;
+  bool causal_mask;
+};
+
+auto& sdpa_cache() {
+  static LRUBytesKeyCache<SDPACacheKey, std::shared_ptr<fe::graph::Graph>>
+      cache(
+          /* capacity */ 128);
+  return cache;
+}
+
+#define Q_UID 1
+#define K_UID 2
+#define V_UID 3
+#define O_UID 4
+#define STATS_UID 5
+
+std::shared_ptr<fe::graph::Graph> get_sdpa_forward_graph(
+    cu::CommandEncoder& encoder,
+    const SDPACacheKey& cache_key) {
+  // Check if graph has already been fully built
+  if (auto it = sdpa_cache().find(cache_key); it != sdpa_cache().end()) {
+    return it->second;
+  }
+
+  // Set up new graph
+  auto graph = std::make_shared<fe::graph::Graph>();
+
+  graph->set_io_data_type(cache_key.cudnn_type)
+      .set_intermediate_data_type(fe::DataType_t::FLOAT)
+      .set_compute_data_type(fe::DataType_t::FLOAT);
+
+  auto Q = graph->tensor(
+      fe::graph::Tensor_attributes()
+          .set_name("Q")
+          .set_uid(Q_UID)
+          .set_dim({cache_key.B, cache_key.H, cache_key.qL, cache_key.D})
+          .set_stride(
+              {cache_key.Q_strides[0],
+               cache_key.Q_strides[1],
+               cache_key.Q_strides[2],
+               1}));
+
+  int h_kv = cache_key.H / cache_key.gqa_factor;
+  auto K =
+      graph->tensor(fe::graph::Tensor_attributes()
+                        .set_name("K")
+                        .set_uid(K_UID)
+                        .set_dim({cache_key.B, h_kv, cache_key.kL, cache_key.D})
+                        .set_stride(
+                            {cache_key.K_strides[0],
+                             cache_key.K_strides[1],
+                             cache_key.V_strides[2],
+                             1}));
+
+  auto V =
+      graph->tensor(fe::graph::Tensor_attributes()
+                        .set_name("V")
+                        .set_uid(V_UID)
+                        .set_dim({cache_key.B, h_kv, cache_key.kL, cache_key.D})
+                        .set_stride(
+                            {cache_key.V_strides[0],
+                             cache_key.V_strides[1],
+                             cache_key.V_strides[2],
+                             1}));
+
+  auto sdpa_options = fe::graph::SDPA_attributes()
+                          .set_name("flash_attention")
+                          .set_is_inference(!cache_key.generate_stats)
+                          .set_attn_scale(cache_key.scale);
+
+  if (cache_key.causal_mask && cache_key.qL > 1) {
+    sdpa_options.set_diagonal_alignment(fe::DiagonalAlignment_t::TOP_LEFT)
+        .set_diagonal_band_right_bound(0);
+  }
+
+  auto [O, Stats] = graph->sdpa(Q, K, V, sdpa_options);
+
+  O->set_output(true)
+      .set_uid(O_UID)
+      .set_dim({cache_key.B, cache_key.H, cache_key.qL, cache_key.D})
+      .set_stride(
+          {cache_key.O_strides[0],
+           cache_key.O_strides[1],
+           cache_key.O_strides[2],
+           1});
+
+  if (cache_key.generate_stats) {
+    Stats->set_output(true)
+        .set_data_type(fe::DataType_t::FLOAT)
+        .set_uid(STATS_UID);
+  }
+
+  // Build and Validate cudnn graph
+
+  auto handle = encoder.device().cudnn_handle();
+
+  // cuDNN only supports native CUDA graphs for sdpa in 9.6 or above.
+  if (cudnnGetVersion() < 90600) {
+    auto build_status = graph->build(handle, {fe::HeurMode_t::A});
+    if (!build_status.is_good()) {
+      throw std::runtime_error(
+          "Unable to build cudnn graph for attention."
+          " Failed with message: " +
+          build_status.get_message());
+    }
+
+  } else {
+    auto val_status = graph->validate();
+    auto op_status = graph->build_operation_graph(handle);
+
+    auto plan_stauts =
+        graph->create_execution_plans({cudnn_frontend::HeurMode_t::A});
+    if (!plan_stauts.is_good()) {
+      throw std::runtime_error(
+          "Unable to create exec plan for cudnn attention."
+          " Failed with message: " +
+          plan_stauts.get_message());
+    }
+
+    graph->select_behavior_notes(
+        {cudnn_frontend::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
+
+    auto support_status = graph->check_support(handle);
+    if (!support_status.is_good()) {
+      throw std::runtime_error(
+          "No cuda graph support for cudnn attention."
+          " Failed with message: " +
+          support_status.get_message());
+    }
+
+    auto build_status = graph->build_plans(handle);
+    if (!build_status.is_good()) {
+      throw std::runtime_error(
+          "Unable to build cudnn graph for attention."
+          " Failed with message: " +
+          build_status.get_message());
+    }
+  }
+
+  auto [it, _] = sdpa_cache().emplace(cache_key, graph);
+
+  return it->second;
+}
+
+inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) {
+  switch (dtype) {
+    case int8:
+      return fe::DataType_t::INT8;
+    case int32:
+      return fe::DataType_t::INT32;
+    case uint8:
+      return fe::DataType_t::UINT8;
+    case float16:
+      return fe::DataType_t::HALF;
+    case bfloat16:
+      return fe::DataType_t::BFLOAT16;
+    case float32:
+      return fe::DataType_t::FLOAT;
+    case float64:
+      return fe::DataType_t::DOUBLE;
+    default:
+      throw std::runtime_error(fmt::format(
+          "Unsupported dtype in SDPA: {}.", dtype_to_string(dtype)));
+  }
+}
+
+void sdpa_cudnn(
+    const Stream& s,
+    cu::CommandEncoder& encoder,
+    const array& q,
+    const array& k,
+    const array& v,
+    const float scale,
+    array& o,
+    bool do_causal_ = false) {
+  encoder.set_input_array(q);
+  encoder.set_input_array(k);
+  encoder.set_input_array(v);
+  encoder.set_output_array(o);
+
+  auto cudnn_type = dtype_to_cudnn_type(q.dtype());
+
+  int B = q.shape(0);
+  int H = q.shape(1);
+  int D = q.shape(3);
+  int gqa_factor = q.shape(1) / k.shape(1);
+
+  int qL = q.shape(2);
+  int kL = k.shape(2);
+
+  SDPACacheKey cache_key{
+      /* int device_id = */ encoder.device().cuda_device(),
+      /* fe::DataType_t cudnn_type = */ cudnn_type,
+
+      /* int B = */ B,
+      /* int H = */ H,
+      /* int D = */ D,
+
+      /* int qL = */ qL,
+      /* int kL = */ kL,
+
+      /* int gqa_factor = */ gqa_factor,
+      /* float scale = */ scale,
+
+      /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)},
+      /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)},
+      /* int64_t V_strides[3] = */ {v.strides(0), v.strides(1), v.strides(2)},
+      /* int64_t O_strides[3] = */ {o.strides(0), o.strides(1), o.strides(2)},
+
+      /* bool generate_stats = */ false,
+      /* bool causal_mask = */ do_causal_};
+
+  auto graph = get_sdpa_forward_graph(encoder, cache_key);
+
+  int64_t workspace_size = 0;
+  auto workspace_status = graph->get_workspace_size(workspace_size);
+  if (!workspace_status.is_good()) {
+    throw std::runtime_error("Unable to get workspace for cudnn attention.");
+  }
+
+  array workspace(
+      allocator::malloc(workspace_size), {int(workspace_size)}, uint8);
+  auto workspace_ptr = workspace.data<void>();
+
+  std::unordered_map<int64_t, void*> variant_pack = {
+      {Q_UID, const_cast<void*>(q.data<void>())},
+      {K_UID, const_cast<void*>(k.data<void>())},
+      {V_UID, const_cast<void*>(v.data<void>())},
+      {O_UID, o.data<void>()}};
+
+  auto handle = encoder.device().cudnn_handle();
+  cudnnSetStream(handle, encoder.stream());
+
+  // cuDNN only supports native CUDA graphs for sdpa in 9.6 or above.
+  if (cudnnGetVersion() < 90600) {
+    auto capture = encoder.capture_context();
+    auto exec_status = graph->execute(handle, variant_pack, workspace_ptr);
+
+    if (!exec_status.is_good()) {
+      capture.discard = true;
+      throw std::runtime_error(
+          "Unable to execute cudnn attention."
+          " Failed with message: " +
+          exec_status.get_message());
+    }
+  } else {
+    cudaGraph_t cu_graph;
+    cudaGraphCreate(&cu_graph, 0);
+
+    std::unique_ptr<cudaGraph_t, void (*)(cudaGraph_t*)> graph_freer(
+        &cu_graph, [](cudaGraph_t* p) { cudaGraphDestroy(*p); });
+
+    auto cu_graph_status = graph->populate_cuda_graph(
+        handle, variant_pack, workspace_ptr, cu_graph);
+
+    if (!cu_graph_status.is_good()) {
+      throw std::runtime_error(
+          "Unable to add cuda graph for cudnn attention."
+          " Failed with message: " +
+          cu_graph_status.get_message());
+    }
+
+    encoder.add_graph_node(cu_graph);
+  }
+
+  encoder.add_temporary(workspace);
+}
+
 } // namespace
 
 namespace fast {
@@ -651,9 +945,6 @@ bool ScaledDotProductAttention::use_fallback(
     bool has_arr_mask,
     bool do_causal,
     Stream s) {
-  if (detail::in_grad_tracing()) {
-    return true;
-  }
   if (s.device == Device::cpu) {
     return true;
   }
@@ -669,7 +960,15 @@ bool ScaledDotProductAttention::use_fallback(
   const bool supported_vector_config =
       sdpa_supported_head_dim && query_sequence_length < 4;
 
-  const bool supported_config = supported_vector_config;
+  auto& cu_device = cu::device(s.device);
+
+  const bool supported_matrix_config = query_sequence_length > 4 &&
+      cu_device.compute_capability_major() >= 8 &&
+      query_sequence_length == key_sequence_length &&
+      (q.dtype() == float16 || q.dtype() == bfloat16);
+
+  const bool supported_config =
+      (supported_matrix_config || supported_vector_config);
 
   return has_arr_mask || !supported_config;
 }
@@ -703,6 +1002,10 @@ void ScaledDotProductAttention::eval_gpu(
     }
   };
 
+  auto is_matrix_contiguous = [](const array& arr) {
+    return arr.strides(-1) == 1;
+  };
+
   // We are in vector mode ie single query
   if (q_pre.shape(2) < 4) {
     auto q_copy_unless = [](const array& arr) {
@@ -756,7 +1059,7 @@ void ScaledDotProductAttention::eval_gpu(
 
       array::Flags flags{
           /* bool contiguous = */ 1,
-          /* bool row_contiguous = */ o.shape(2) == 1,
+          /* bool row_contiguous = */ 0,
           /* bool col_contiguous = */ 0,
       };
 
@@ -770,9 +1073,35 @@ void ScaledDotProductAttention::eval_gpu(
     return sdpa_vector_fallback(s, encoder, q, k, v, scale_, o, do_causal_);
   }
 
-  // Full attention mode should never reach here
+  // Full attention mode
   else {
-    throw std::runtime_error("Doesn't support matrix yet.");
+    const auto& q = copy_unless(is_matrix_contiguous, q_pre);
+    const auto& k = copy_unless(is_matrix_contiguous, k_pre);
+    const auto& v = copy_unless(is_matrix_contiguous, v_pre);
+
+    for (const auto& cp : copies) {
+      encoder.add_temporary(cp);
+    }
+
+    int64_t str_oD = 1;
+    int64_t str_oH = o.shape(3);
+    int64_t str_oL = o.shape(1) * str_oH;
+    int64_t str_oB = o.shape(2) * str_oL;
+    size_t data_size = o.shape(0) * str_oB;
+
+    array::Flags flags{
+        /* bool contiguous = */ 1,
+        /* bool row_contiguous = */ 0,
+        /* bool col_contiguous = */ 0,
+    };
+
+    o.set_data(
+        allocator::malloc(o.nbytes()),
+        data_size,
+        {str_oB, str_oH, str_oL, str_oD},
+        flags);
+
+    return sdpa_cudnn(s, encoder, q, k, v, scale_, o, do_causal_);
   }
 }
 

mlx/backend/metal/device.h

@@ -104,7 +104,7 @@ struct CommandEncoder {
   };
 
   // Outputs of all kernels in the encoder including temporaries
-  std::unordered_set<const void*>& outputs() {
+  std::unordered_set<const void*> outputs() {
     return all_outputs_;
   };
 

mlx/version.h

@@ -3,8 +3,8 @@
 #pragma once
 
 #define MLX_VERSION_MAJOR 0
-#define MLX_VERSION_MINOR 28
-#define MLX_VERSION_PATCH 0
+#define MLX_VERSION_MINOR 27
+#define MLX_VERSION_PATCH 1
 #define MLX_VERSION_NUMERIC \
   (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH)
 

python/mlx/nn/layers/base.py

@@ -178,7 +178,7 @@ class Module(dict):
 
         if strict:
             new_weights = dict(weights)
-            curr_weights = tree_flatten(self.parameters(), destination={})
+            curr_weights = dict(tree_flatten(self.parameters()))
             if extras := (new_weights.keys() - curr_weights.keys()):
                 num_extra = len(extras)
                 extras = ",\n".join(sorted(extras))
@@ -212,7 +212,7 @@ class Module(dict):
         - ``.npz`` will use :func:`mx.savez`
         - ``.safetensors`` will use :func:`mx.save_safetensors`
         """
-        params_dict = tree_flatten(self.parameters(), destination={})
+        params_dict = dict(tree_flatten(self.parameters()))
 
         if file.endswith(".npz"):
             mx.savez(file, **params_dict)

python/mlx/utils.py

@@ -1,7 +1,7 @@
 # Copyright © 2023 Apple Inc.
 from collections import defaultdict
 from itertools import zip_longest
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, List, Optional, Tuple
 
 
 def tree_map(
@@ -114,11 +114,8 @@ def tree_map_with_path(
 
 
 def tree_flatten(
-    tree: Any,
-    prefix: str = "",
-    is_leaf: Optional[Callable] = None,
-    destination: Optional[Union[List[Tuple[str, Any]], Dict[str, Any]]] = None,
-) -> Union[List[Tuple[str, Any]], Dict[str, Any]]:
+    tree: Any, prefix: str = "", is_leaf: Optional[Callable] = None
+) -> Any:
     """Flattens a Python tree to a list of key, value tuples.
 
     The keys are using the dot notation to define trees of arbitrary depth and
@@ -131,12 +128,9 @@ def tree_flatten(
         print(tree_flatten([[[0]]]))
         # [("0.0.0", 0)]
 
-        print(tree_flatten([[[0]]], prefix=".hello"))
+        print(tree_flatten([[[0]]], ".hello"))
         # [("hello.0.0.0", 0)]
 
-        tree_flatten({"a": {"b": 1}}, destination={})
-        {"a.b": 1}
-
     .. note::
        Dictionaries should have keys that are valid Python identifiers.
 
@@ -146,50 +140,26 @@ def tree_flatten(
             always discarded.
         is_leaf (callable): An optional callable that returns True if the
             passed object is considered a leaf or False otherwise.
-        destination (list or dict, optional): A list or dictionary to store the
-            flattened tree. If None an empty list will be used. Default: ``None``.
 
     Returns:
-        Union[List[Tuple[str, Any]], Dict[str, Any]]: The flat representation of
-            the Python tree.
+        List[Tuple[str, Any]]: The flat representation of the Python tree.
     """
-    if destination is None:
-        destination = []
+    flat_tree = []
 
-    # Create the function to update the destination. We are taking advantage of
-    # the fact that list.extend and dict.update have the same API to simplify
-    # the code a bit.
-    if isinstance(destination, list):
-        _add_to_destination = destination.extend
-    elif isinstance(destination, dict):
-        _add_to_destination = destination.update
-    else:
-        raise ValueError("Destination should be either a list or a dictionary or None")
+    if is_leaf is None or not is_leaf(tree):
+        if isinstance(tree, (list, tuple)):
+            for i, t in enumerate(tree):
+                flat_tree.extend(tree_flatten(t, f"{prefix}.{i}", is_leaf))
+            return flat_tree
+        if isinstance(tree, dict):
+            for k, t in tree.items():
+                flat_tree.extend(tree_flatten(t, f"{prefix}.{k}", is_leaf))
+            return flat_tree
 
-    # Leaf identified by is_leaf so add it and return
-    if is_leaf is not None and is_leaf(tree):
-        _add_to_destination([(prefix[1:], tree)])
-        return destination
-
-    # List or tuple so recursively add each subtree
-    if isinstance(tree, (list, tuple)):
-        for i, item in enumerate(tree):
-            tree_flatten(item, f"{prefix}.{i}", is_leaf, destination)
-        return destination
-
-    # Dictionary so recursively add each subtree
-    if isinstance(tree, dict):
-        for key, value in tree.items():
-            tree_flatten(value, f"{prefix}.{key}", is_leaf, destination)
-        return destination
-
-    # Leaf so add it and return
-    _add_to_destination([(prefix[1:], tree)])
-
-    return destination
+    return [(prefix[1:], tree)]
 
 
-def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
+def tree_unflatten(tree: List[Tuple[str, Any]]) -> Any:
     """Recreate a Python tree from its flat representation.
 
     .. code-block:: python
@@ -200,34 +170,31 @@ def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
         print(d)
         # {"hello": {"world": 42}}
 
-        d = tree_unflatten({"hello.world": 42})
-        print(d)
-        # {"hello": {"world": 42}}
-
     Args:
-        tree (list[tuple[str, Any]] or dict[str, Any]): The flat representation of a Python tree.
+        tree (list[tuple[str, Any]]): The flat representation of a Python tree.
            For instance as returned by :meth:`tree_flatten`.
 
     Returns:
         A Python tree.
     """
-    items = tree.items() if isinstance(tree, dict) else tree
+    if len(tree) == 1 and tree[0][0] == "":
+        return tree[0][1]
 
-    # Special case when we have just one element in the tree ie not a tree
-    if len(items) == 1:
-        key, value = next(iter(items))
-        if key == "":
-            return value
+    try:
+        int(tree[0][0].split(".", maxsplit=1)[0])
+        is_list = True
+    except ValueError:
+        is_list = False
 
     # collect children
     children = defaultdict(list)
-    for key, value in items:
+    for key, value in tree:
         current_idx, *next_idx = key.split(".", maxsplit=1)
         next_idx = "" if not next_idx else next_idx[0]
         children[current_idx].append((next_idx, value))
 
-    # Assume they are a list and fail to dict if the keys are not all integers
-    try:
+    # recursively map them to the original container
+    if is_list:
         keys = sorted((int(idx), idx) for idx in children.keys())
         l = []
         for i, k in keys:
@@ -235,7 +202,7 @@ def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
             l.extend([{} for _ in range(i - len(l))])
             l.append(tree_unflatten(children[k]))
         return l
-    except ValueError:
+    else:
         return {k: tree_unflatten(v) for k, v in children.items()}
 
 

python/tests/test_nn.py

@@ -80,7 +80,7 @@ class TestBase(mlx_tests.MLXTestCase):
                 self.weights = {"w1": mx.zeros((2, 2)), "w2": mx.ones((2, 2))}
 
         model = DictModule()
-        params = tree_flatten(model.parameters(), destination={})
+        params = dict(tree_flatten(model.parameters()))
         self.assertEqual(len(params), 2)
         self.assertTrue(mx.array_equal(params["weights.w1"], mx.zeros((2, 2))))
         self.assertTrue(mx.array_equal(params["weights.w2"], mx.ones((2, 2))))