Allow some non 2D inputs in qqmm (#2981)

Co-authored-by: KD2YCU <me@kd2ycu.com>
Co-authored-by: Angelos Katharopoulos <katharas@gmail.com>
Co-authored-by: Awni Hannun <awni@apple.com>

.github/actions/build-macos/action.yml

@@ -4,12 +4,14 @@ description: 'Build and test MLX on macOS'
 runs:
   using: "composite"
   steps:
-    - name: Build
+    - name: Install dependencies
       env:
         DEBUG: 1
         CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON"
       shell: bash -l {0}
       run: |
+        pip install --upgrade pip
+        pip install cmake setuptools typing_extensions
         pip install -e . -v
 
     - name: Install tests dependencies

.github/actions/setup-macos/action.yml

@@ -18,18 +18,7 @@ runs:
       shell: bash
       run: xcodebuild -showComponent MetalToolchain
 
-    - name: Install Python
-      shell: sh
-      run: |
-        curl -LsSf https://astral.sh/uv/install.sh | sh
-        $HOME/.local/bin/uv venv --python ${{ inputs.python-version }}
-        source .venv/bin/activate
-        echo PATH=$PATH >> $GITHUB_ENV
-        # Search python packages in .venv
-        echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV
-
-    - name: Install build dependencies
-      shell: sh
-      run: |
-        pip install --upgrade pip
-        pip install cmake setuptools typing_extensions
+    - uses: conda-incubator/setup-miniconda@v3
+      with:
+        miniconda-version: "latest"
+        python-version: ${{ inputs.python-version }}

docs/src/usage/distributed.rst

@@ -584,7 +584,7 @@ the process.
 **MLX_JACCL_COORDINATOR** should contain the IP and port that rank 0 can listen
 to all the other ranks connect to in order to establish the RDMA connections.
 
-**MLX_IBV_DEVICES** should contain the path to a json file that contains the
+**MLX_JACCL_DEVICES** should contain the path to a json file that contains the
 ibverbs device names that connect each node to each other node, something like
 the following:
 

mlx/backend/cpu/unary_ops.h

@@ -154,24 +154,12 @@ struct ToFP8 {
 struct FromFP8 {
   template <int N>
   Simd<float, N> operator()(Simd<uint8_t, N> x) {
-    auto w = Simd<uint32_t, N>(x) << 24;
-    auto sign = w & 0x80000000;
-    auto nonsign = w & 0x7FFFFFFF;
-
-    auto renorm_shift = clz(nonsign);
-    renorm_shift = simd::select(
-        renorm_shift > Simd<uint32_t, N>{4},
-        renorm_shift - Simd<uint32_t, N>{4},
-        Simd<uint32_t, N>{0});
-
-    Simd<int32_t, N> inf_nan_mask =
-        (Simd<int32_t, N>(nonsign + 0x01000000) >> 8) & 0x7F800000;
-    auto zero_mask = Simd<int32_t, N>(nonsign - 1) >> 31;
-    auto result = sign |
-        ((((nonsign << renorm_shift >> 4) + ((0x78 - renorm_shift) << 23)) |
-          inf_nan_mask) &
-         ~zero_mask);
-    return fp32_from_bits(result);
+    auto v = Simd<uint16_t, N>(x & 127) << 7;
+    auto converted = *(Simd<float16_t, N>*)(&v);
+    converted = converted * 256.0;
+    auto sign = Simd<bool, N>(x & 128);
+    Simd<float, N> out = select(sign, -converted, converted);
+    return out;
   }
   float operator()(uint8_t x) {
     return (*this)(Simd<uint8_t, 1>(x)).value;

mlx/backend/cuda/quantized/affine_quantize.cu

@@ -23,8 +23,7 @@ affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) {
   auto tidx = block_idx.x * block_size.x + idx_in_block.x;
   auto tidy = block_idx.y * block_size.y + idx_in_block.y;
 
-  auto grid_dim_x =
-      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
   constexpr float eps = 1e-7;
   constexpr int simd_size = WARP_SIZE;
   constexpr float n_bins = (1 << bits) - 1;
@@ -141,8 +140,7 @@ __global__ void affine_dequantize(
   auto tidx = block_idx.x * block_size.x + idx_in_block.x;
   auto tidy = block_idx.y * block_size.y + idx_in_block.y;
 
-  auto grid_dim_x =
-      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
+  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;
 
   constexpr int pack_factor = get_pack_factor<bits, 8>();
   constexpr int bytes_per_pack = get_bytes_per_pack<bits>();

mlx/backend/cuda/quantized/qqmm.cpp

@@ -46,7 +46,7 @@ inline array ensure_row_contiguous_matrix(
   return x_copy;
 }
 
-array pad_and_repack_scales(
+array pad_and_swizzle_scales(
     const array& scale,
     cu::CommandEncoder& encoder,
     const Stream& s) {
@@ -64,14 +64,12 @@ array pad_and_repack_scales(
       cu::malloc_async(pad_outer * pad_inner, encoder),
       Shape{pad_outer, pad_inner},
       scale.dtype());
-  repack_scales(scale, scale_tiled, encoder, s);
+  swizzle_scales(scale, scale_tiled, encoder, s);
 
   encoder.add_temporary(scale_tiled);
   return scale_tiled;
 }
-} // namespace
 
-namespace {
 void qqmm_impl(
     cu::CommandEncoder& encoder,
     int M,
@@ -176,8 +174,8 @@ void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   int K = K_packed * (32 / bits_);
 
   // Repack scales from linear to tiled layout for tensor cores
-  array scale_x = pad_and_repack_scales(scale_x_pre, encoder, s);
-  array scale_w = pad_and_repack_scales(scale_w_pre, encoder, s);
+  array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s);
+  array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s);
 
   bool x_transposed = false;
   bool w_transposed = true; // always transposed

mlx/backend/cuda/quantized/qqmm_utils.cu

@@ -10,6 +10,11 @@ namespace mlx::core {
 
 namespace cg = cooperative_groups;
 
+constexpr int TILE_ROWS = 128;
+constexpr int TILE_COLS = 4;
+constexpr int TILES_PER_LANE = 1;
+constexpr int LANES_PER_BLOCK = 32;
+
 // To pass scales to tensor cores, they need to be repacked into a tiled layout
 // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
 // Tiled layout for scale factors is very well described in CUTLASS
@@ -43,118 +48,171 @@ namespace cg = cooperative_groups;
 //          [252, 253, 254, 255],
 //          [380, 381, 382, 383],
 //          [508, 509, 510, 511]]]]],
-__device__ size_t
-scale_tiled_offset(size_t scale_index, size_t num_rows, size_t num_scale_cols) {
-  // Compute the tiled layout offset for scale factors used in tensor cores
-  // This function maps from a linear scale index to the tiled layout expected
-  // by tensor cores (and cublaslt).
-  //
-  // Input: linear scale index (e.g., for a matrix M x K with group_size,
-  //        scale_index ranges from 0 to (M * K/group_size - 1))
-  //
-  // The tiled layout organizes scales into tiles of 128 rows x 4 columns,
-  // where each tile is subdivided into 4 sub-blocks of 32 rows x 4 columns.
-  size_t row = scale_index / num_scale_cols;
-  size_t col = scale_index % num_scale_cols;
 
-  constexpr size_t rows_per_tile = 128;
-  constexpr size_t rows_per_sub_block = 32;
-  constexpr size_t cols_per_sub_block = 4;
-  constexpr size_t sub_blocks_per_tile = 4; // Vertically stacked
+inline std::tuple<dim3, dim3> get_swizzle_launch_args(
+    size_t M_swizzled,
+    size_t K_swizzled) {
+  constexpr int tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
+  constexpr int warps_per_block = TILE_ROWS / 4; // 128 / 4 = 32
 
-  // Decompose row position
-  size_t tile_row = row / rows_per_tile; // Which tile row
-  size_t row_in_tile = row % rows_per_tile; // Row within tile
-  size_t sub_block_row =
-      row_in_tile / rows_per_sub_block; // Sub-block within tile
-  size_t row_in_sub_block =
-      row_in_tile % rows_per_sub_block; // Row in sub-block
+  const int num_tiles_k = K_swizzled / TILE_COLS;
+  const int num_tiles_m = M_swizzled / TILE_ROWS;
 
-  // Decompose column position
-  size_t col_tile = col / cols_per_sub_block; // Which column tile
-  size_t col_in_sub_block = col % cols_per_sub_block; // Column within sub-block
+  dim3 grid;
+  grid.x = cuda::ceil_div(num_tiles_k, tiles_per_block);
+  grid.y = num_tiles_m;
+  grid.z = 1;
+  // Block is always (32, 32) = 1024 threads
+  dim3 block(LANES_PER_BLOCK, warps_per_block, 1);
 
-  // Compute tile index and offset within tile
-  size_t num_col_tiles = cuda::ceil_div(num_scale_cols, cols_per_sub_block);
-  size_t tile_idx = tile_row * num_col_tiles + col_tile;
-
-  size_t offset_in_tile =
-      (row_in_sub_block * sub_blocks_per_tile * cols_per_sub_block) +
-      (sub_block_row * cols_per_sub_block) + col_in_sub_block;
-
-  constexpr size_t tile_size = rows_per_tile * cols_per_sub_block;
-  return tile_idx * tile_size + offset_in_tile;
+  return std::make_tuple(grid, block);
 }
 
 namespace cu {
 
-__global__ void repack_scales(
+__global__ void swizzle_scales(
     const uint8_t* scales_linear,
-    uint8_t* scales_tiled,
-    size_t input_rows,
-    size_t input_cols,
-    size_t output_rows,
-    size_t output_cols) {
+    uint8_t* scales_swizzled,
+    const size_t M,
+    const size_t K,
+    const size_t M_swizzled,
+    const size_t K_swizzled) {
+  constexpr int tile_size = TILE_ROWS * TILE_COLS;
+  constexpr int num_tile_rows_per_thread = 4;
+  constexpr int max_tiles_per_block = LANES_PER_BLOCK * TILES_PER_LANE;
+
+  constexpr int tile_stride = tile_size / 16; // 32 int4s per tile
+
+  // Each thread loads 16 scales from 4 rows (stride 32) and packs them into
+  // int4. For example: thread (0, 0) loads scales at rows 0,32,64,96 of tile 0,
+  // thread (1, 0) loads rows 0,32,64,96 of of tile 1, etc.
+  // The store is strided within a warp (stride 32 int4s), so we first
+  // write to shared memory, then do a coalesced store from shared to global
   auto block_size = cg::this_thread_block().dim_threads();
   auto block_idx = cg::this_thread_block().group_index();
   auto idx_in_block = cg::this_thread_block().thread_index();
 
-  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
-  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
+  auto tidx = idx_in_block.x;
+  auto tidy = idx_in_block.y;
+  auto linear_tid = tidy * block_size.x + tidx;
 
-  auto grid_dim_x =
-      cg::this_grid().dim_blocks().x * cg::this_grid().block_index().x;
+  const int bid_x = block_idx.x;
+  const int bid_y = block_idx.y;
 
-  size_t output_index = tidx + grid_dim_x * size_t(tidy);
-  size_t output_size = output_rows * output_cols;
+  const int K_int = K_swizzled / 4;
 
-  if (output_index >= output_size) {
-    return;
+  const size_t output_offset = static_cast<size_t>(bid_y) * TILE_ROWS * K_int +
+      static_cast<size_t>(bid_x) * max_tiles_per_block * tile_size / 4;
+  int* output_block = reinterpret_cast<int*>(scales_swizzled) + output_offset;
+
+  const int grid_dim_x = cg::this_grid().dim_blocks().x;
+  const int grid_dim_y = cg::this_grid().dim_blocks().y;
+
+  int remaining = K_int - bid_x * max_tiles_per_block;
+  int tiles_in_block = min(remaining, max_tiles_per_block);
+  bool valid_tile = tidx * TILES_PER_LANE < tiles_in_block;
+
+  __shared__ int4 strided_scales_thread[max_tiles_per_block * tile_stride];
+
+  // Initialize to zero for padding
+  int thread_tile_rows[num_tile_rows_per_thread] = {0};
+
+  if (valid_tile) {
+    const size_t col_base =
+        static_cast<size_t>(bid_x) * max_tiles_per_block * TILE_COLS +
+        tidx * TILE_COLS;
+
+    const bool aligned_k = (K % 4 == 0);
+
+    if (aligned_k) {
+      // fast path: K is aligned, use vectorized loads with stride K/4
+      const int K_stride = K / 4;
+      const size_t block_offset =
+          static_cast<size_t>(bid_y) * TILE_ROWS * K_stride +
+          static_cast<size_t>(bid_x) * max_tiles_per_block;
+      const int* input_block =
+          reinterpret_cast<const int*>(scales_linear) + block_offset;
+// load
+#pragma unroll
+      for (int i = 0; i < num_tile_rows_per_thread; i++) {
+        const size_t row =
+            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
+        const int thread_offset =
+            (i * block_size.x + tidy) * K_stride + tidx * TILES_PER_LANE;
+        if (row < M && col_base + TILE_COLS <= K) {
+          thread_tile_rows[i] = __ldg(input_block + thread_offset);
+        } else if (row < M) {
+// partial tile at K boundary: load byte-by-byte
+#pragma unroll
+          for (int c = 0; c < TILE_COLS; c++) {
+            if (col_base + c < K) {
+              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
+                  scales_linear[row * K + col_base + c];
+            }
+          }
+        }
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < num_tile_rows_per_thread; i++) {
+        const size_t row =
+            static_cast<size_t>(bid_y) * TILE_ROWS + i * block_size.x + tidy;
+        if (row < M) {
+          const size_t row_start = row * K;
+#pragma unroll
+          for (int c = 0; c < TILE_COLS; c++) {
+            if (col_base + c < K) {
+              reinterpret_cast<uint8_t*>(&thread_tile_rows[i])[c] =
+                  scales_linear[row_start + col_base + c];
+            }
+          }
+        }
+      }
+    }
+    // store to shared with XOR swizzle to avoid bank conflicts
+    int base_idx = tidx * tile_stride + tidy;
+    int xor_bits = (tidy >> 3) & 0x3;
+    int swizzled_idx = base_idx ^ xor_bits;
+    strided_scales_thread[swizzled_idx] =
+        *reinterpret_cast<int4*>(thread_tile_rows);
   }
 
-  size_t tiled_offset =
-      scale_tiled_offset(output_index, output_rows, output_cols);
+  cg::thread_block block = cg::this_thread_block();
+  cg::sync(block);
 
-  size_t row = output_index / output_cols;
-  size_t col = output_index % output_cols;
-
-  // Probably this can be done better with 2 separated paths for valid and
-  // padding
-  if (row < input_rows && col < input_cols) {
-    size_t input_index = row * input_cols + col;
-    scales_tiled[tiled_offset] = scales_linear[input_index];
-  } else {
-    // Zero-fill padding region
-    scales_tiled[tiled_offset] = 0;
+  const int total_int4s = tiles_in_block * tile_stride;
+#pragma unroll
+  for (int i = linear_tid; i < total_int4s; i += block_size.x * block_size.y) {
+    int tile_idx = i / tile_stride;
+    int row_idx = i % tile_stride;
+    int base_idx = tile_idx * tile_stride + row_idx;
+    int xor_bits = (row_idx >> 3) & 0x3;
+    int swizzled_idx = base_idx ^ xor_bits;
+    reinterpret_cast<int4*>(output_block)[i] =
+        strided_scales_thread[swizzled_idx];
   }
 }
-
 } // namespace cu
 
-void repack_scales(
+void swizzle_scales(
     const array& scales,
     array& scales_tiled,
     cu::CommandEncoder& enc,
     const Stream& s) {
   enc.set_input_array(scales);
   enc.set_output_array(scales_tiled);
-
   // Note: scales_tiled is padded to full tiles so if num_rows or num_cols
-  // are not multiples of tile sizes, the extra space is filled with zeros
-
+  // are not multiples of tile sizes
   size_t input_rows = scales.shape(-2);
   size_t input_cols = scales.shape(-1);
 
   size_t output_rows = scales_tiled.shape(-2);
   size_t output_cols = scales_tiled.shape(-1);
-  size_t output_size = output_rows * output_cols;
-
-  bool large = output_size > UINT_MAX;
-  auto [num_blocks, block_dims] = get_launch_args(
-      output_size, scales_tiled.shape(), scales_tiled.strides(), large);
 
+  auto [num_blocks, block_dims] =
+      get_swizzle_launch_args(output_rows, output_cols);
   enc.add_kernel_node(
-      cu::repack_scales,
+      cu::swizzle_scales,
       num_blocks,
       block_dims,
       0,

mlx/backend/cuda/quantized/qqmm_utils.h

@@ -21,7 +21,7 @@ inline std::pair<int, int> get_padded_scale_dims(int num_rows, int num_cols) {
   return {padded_rows, padded_cols};
 }
 
-void repack_scales(
+void swizzle_scales(
     const array& scales,
     array& scales_tiled,
     cu::CommandEncoder& enc,

mlx/backend/metal/device.h

@@ -266,6 +266,9 @@ Device& device(mlx::core::Device);
 std::unique_ptr<void, std::function<void(void*)>> new_scoped_memory_pool();
 
 inline bool is_nax_available() {
+#ifdef MLX_METAL_NO_NAX
+  return false;
+#else
   auto _check_nax = []() {
     bool can_use_nax = false;
     if (__builtin_available(
@@ -278,6 +281,7 @@ inline bool is_nax_available() {
   };
   static bool is_nax_available_ = _check_nax();
   return is_nax_available_;
+#endif
 }
 
 } // namespace mlx::core::metal

mlx/backend/metal/kernels/CMakeLists.txt

@@ -164,6 +164,8 @@ if(NOT MLX_METAL_JIT)
     build_kernel(steel/attn/kernels/steel_attention_nax
                  ${STEEL_NAX_ATTN_HEADERS})
 
+  else()
+    target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX)
   endif()
 
 endif()

mlx/distributed/jaccl/jaccl.cpp

@@ -1159,7 +1159,7 @@ bool is_available() {
 }
 
 std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
-  const char* dev_file = std::getenv("MLX_IBV_DEVICES");
+  const char* dev_file = std::getenv("MLX_JACCL_DEVICES");
   const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR");
   const char* rank_str = std::getenv("MLX_RANK");
 
@@ -1167,9 +1167,9 @@ std::shared_ptr<GroupImpl> init(bool strict /* = false */) {
     if (strict) {
       std::ostringstream msg;
       msg << "[jaccl] You need to provide via environment variables a rank (MLX_RANK), "
-          << "a device file (MLX_IBV_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) "
+          << "a device file (MLX_JACCL_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) "
           << "but provided MLX_RANK=\"" << ((rank_str) ? rank_str : "")
-          << "\", MLX_IBV_DEVICES=\"" << ((dev_file) ? dev_file : "")
+          << "\", MLX_JACCL_DEVICES=\"" << ((dev_file) ? dev_file : "")
           << "\" and MLX_JACCL_COORDINATOR=\""
           << ((coordinator) ? coordinator : "");
       throw std::runtime_error(msg.str());

mlx/io/safetensors.cpp

@@ -95,7 +95,6 @@ Dtype dtype_from_safetensor_str(std::string_view str) {
   } else if (str == ST_C64) {
     return complex64;
   } else if (str == ST_F8_E4M3) {
-    // We convert this manually later
     return uint8;
   } else {
     throw std::runtime_error(
@@ -148,16 +147,14 @@ SafetensorsLoad load_safetensors(
     const Shape& shape = item.value().at("shape");
     const std::vector<size_t>& data_offsets = item.value().at("data_offsets");
     Dtype type = dtype_from_safetensor_str(dtype);
-    auto loaded_array = array(
-        shape,
-        type,
-        std::make_shared<Load>(
-            stream, in_stream, offset + data_offsets.at(0), false),
-        std::vector<array>{});
-    if (dtype == ST_F8_E4M3) {
-      loaded_array = from_fp8(loaded_array, bfloat16, s);
-    }
-    res.insert({item.key(), loaded_array});
+    res.insert(
+        {item.key(),
+         array(
+             shape,
+             type,
+             std::make_shared<Load>(
+                 stream, in_stream, offset + data_offsets.at(0), false),
+             std::vector<array>{})});
   }
   return {res, metadata_map};
 }

mlx/ops.cpp

@@ -4335,7 +4335,7 @@ std::pair<int, int> extract_qqmm_dims(
 }
 
 array qqmm(
-    array x,
+    array in_x,
     array w,
     std::optional<array> scales_w,
     std::optional<int> group_size_ /* = std::nullopt */,
@@ -4360,6 +4360,16 @@ array qqmm(
   // 2. w is not quantized, scales is not provided
   auto [group_size, bits] =
       quantization_params_from_mode(qmode, group_size_, bits_);
+
+  // Allow gemv
+  auto x = in_x;
+  if (x.ndim() == 1) {
+    // Insert a singleton dim in the beginning
+    x = expand_dims(x, 0, s);
+  } else if (w.ndim() == 2 && x.ndim() > 2) {
+    x = flatten(x, 0, -2, s);
+  }
+
   // validate inputs
   validate_qqmm_inputs(x, w, scales_w, group_size, bits);
   // validate and extract shapes
@@ -4374,11 +4384,19 @@ array qqmm(
   }
   auto out_shape = inputs[0].shape();
   out_shape.back() = w_outer_dims;
-  return array(
+  auto out = array(
       std::move(out_shape),
       x.dtype(), // output dtype is the same as x dtype
       std::make_shared<QQMatmul>(stream, group_size, bits, qmode),
       std::move(inputs));
+  if (in_x.ndim() > 2) {
+    auto orig_shape = in_x.shape();
+    orig_shape.pop_back();
+    out = unflatten(out, 0, std::move(orig_shape), s);
+  } else if (in_x.ndim() == 1) {
+    out = squeeze(out, 0, s);
+  }
+  return out;
 }
 
 array pack_and_quantize(

mlx/primitives.cpp

@@ -3717,7 +3717,7 @@ std::pair<std::vector<array>, std::vector<int>> RandomBits::vmap(
 
 bool RandomBits::is_equivalent(const Primitive& other) const {
   const RandomBits& r_other = static_cast<const RandomBits&>(other);
-  return shape_ == r_other.shape_;
+  return shape_ == r_other.shape_ && width_ == r_other.width_;
 }
 
 std::vector<array> Real::vjp(

python/mlx/_distributed_utils/launch.py

@@ -376,7 +376,7 @@ def launch_jaccl(parser, hosts, args, command):
     env = args.env
     cwd = args.cwd
     env.append(f"MLX_JACCL_COORDINATOR={coordinator}:{args.starting_port}")
-    files = {"MLX_IBV_DEVICES": json.dumps([h.rdma for h in hosts])}
+    files = {"MLX_JACCL_DEVICES": json.dumps([h.rdma for h in hosts])}
 
     log(args.verbose, "Running", shlex.join(command))
 

python/src/ops.cpp

@@ -1668,6 +1668,25 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The array of zeros with the specified shape.
       )pbdoc");
+  m.def(
+      "asarray",
+      [](const ArrayInitType& a, std::optional<mx::Dtype> dtype) {
+        return create_array(a, dtype);
+      },
+      nb::arg(),
+      "dtype"_a = nb::none(),
+      nb::sig("def asarray(a: Union[scalar, array, Sequence], dtype: "
+              "Optional[Dtype] = None) -> array"),
+      R"pbdoc(
+        Convert the input to an array.
+
+        Args:
+            a: Input data.
+            dtype (Dtype, optional): The desired data-type for the array.
+
+        Returns:
+            array: An array interpretation of the input.
+      )pbdoc");
   m.def(
       "zeros_like",
       &mx::zeros_like,
@@ -4281,7 +4300,7 @@ void init_ops(nb::module_& m) {
             ======  ======================   ==========================  =============  =====
             affine  32, 64\ :sup:`*`, 128    2, 3, 4\ :sup:`*`, 5, 6, 8  same as input  yes
             mxfp4   32\ :sup:`*`             4\ :sup:`*`                 e8m0           no
-            mxfp8   32\ :sup:`*`             4\ :sup:`*`                 e8m0           no
+            mxfp8   32\ :sup:`*`             8\ :sup:`*`                 e8m0           no
             nvfp4   16\ :sup:`*`             4\ :sup:`*`                 e4m3           no
             ======  ======================   ==========================  =============  =====
 
@@ -4316,7 +4335,7 @@ void init_ops(nb::module_& m) {
           size must be 16. The elements are quantized to 4-bit or 8-bit
           precision floating-point values: E2M1 for ``"fp4"`` and E4M3 for
           ``"fp8"``. There is a shared 8-bit scale per group. The ``"mx"``
-          modes us an E8M0 scale and the ``"nv"`` mode uses an E4M3 scale.
+          modes use an E8M0 scale and the ``"nv"`` mode uses an E4M3 scale.
           Unlike ``affine`` quantization, these modes does not have a bias
           value.
 
@@ -5460,10 +5479,10 @@ void init_ops(nb::module_& m) {
         If ``w`` is expected to receive gradients, it must be provided in
         non-quantized form.
 
-        If ``x`` and `w`` are not quantized, their data types must be ``float32``, 
+        If ``x`` and `w`` are not quantized, their data types must be ``float32``,
         ``float16``, or ``bfloat16``.
         If ``w`` is quantized, it must be packed in unsigned integers.
-        
+
       Args:
         x (array): Input array.
         w (array): Weight matrix. If quantized, it is packed in unsigned integers.
@@ -5483,4 +5502,40 @@ void init_ops(nb::module_& m) {
         array: The result of the multiplication of quantized ``x`` with quantized ``w``.
         needed).
   )pbdoc");
+  m.def(
+      "from_fp8",
+      &mx::from_fp8,
+      nb::arg(),
+      "dtype"_a = mx::bfloat16,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+      Convert the array from fp8 (e4m3) to another floating-point type.
+
+      Args:
+        x (array): The input fp8 array with type ``uint8``.
+        dtype (Dtype): The data type to convert to. Default: ``bfloat16``.
+
+      Returns:
+        array: The array converted from fp8.
+  )pbdoc");
+  m.def(
+      "to_fp8",
+      &mx::to_fp8,
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def to_fp8(x: array, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+      Convert the array to fp8 (e4m3) from another floating-point type.
+
+      Args:
+        x (array): The input array.
+
+      Returns:
+        array: The array converted to fp8 with type ``uint8``.
+  )pbdoc");
 }

python/tests/cuda_skip.py

@@ -1,5 +1,4 @@
 cuda_skip = {
-    "TestLoad.test_load_f8_e4m3",
     "TestLayers.test_quantized_embedding",
     # Block masked matmul NYI
     "TestBlas.test_block_masked_matmul",

python/tests/test_array.py

@@ -1973,6 +1973,58 @@ class TestArray(mlx_tests.MLXTestCase):
         self.assertTrue(hasattr(api, "array"))
         self.assertTrue(hasattr(api, "add"))
 
+    def test_array_namespace_asarray(self):
+        xp = mx.array(1.0).__array_namespace__()
+        self.assertTrue(hasattr(xp, "asarray"))
+
+        arr = xp.asarray([1, 2, 3])
+        self.assertEqual(arr.tolist(), [1, 2, 3])
+
+        arr_f32 = xp.asarray([1, 2, 3], dtype=mx.float32)
+        self.assertEqual(arr_f32.dtype, mx.float32)
+
+        existing = mx.array([4, 5, 6])
+        arr_pass = xp.asarray(existing)
+        self.assertEqual(arr_pass.tolist(), [4, 5, 6])
+
+    def test_asarray(self):
+        # List inputs
+        self.assertEqual(mx.asarray([1, 2, 3]).tolist(), [1, 2, 3])
+        self.assertEqual(mx.asarray([[1, 2], [3, 4]]).tolist(), [[1, 2], [3, 4]])
+
+        # Tuple inputs
+        self.assertEqual(mx.asarray((1, 2, 3)).tolist(), [1, 2, 3])
+        self.assertEqual(mx.asarray(((1, 2), (3, 4))).tolist(), [[1, 2], [3, 4]])
+
+        # Mixed nesting
+        self.assertEqual(mx.asarray([(1, 2), (3, 4)]).tolist(), [[1, 2], [3, 4]])
+        self.assertEqual(mx.asarray(([1, 2], [3, 4])).tolist(), [[1, 2], [3, 4]])
+
+        # Scalar inputs
+        self.assertEqual(mx.asarray(42).item(), 42)
+        self.assertEqual(mx.asarray(3.14).item(), 3.140000104904175)
+        self.assertEqual(mx.asarray(True).item(), True)
+        self.assertEqual(mx.asarray(1 + 2j).item(), (1 + 2j))
+
+        # MLX array inputs
+        arr = mx.array([1, 2, 3])
+        self.assertEqual(mx.asarray(arr).tolist(), [1, 2, 3])
+
+        arr_int = mx.array([1, 2, 3], dtype=mx.int32)
+        arr_float = mx.asarray(arr_int, dtype=mx.float32)
+        self.assertEqual(arr_float.dtype, mx.float32)
+        self.assertEqual(arr_float.tolist(), [1.0, 2.0, 3.0])
+
+        # NumPy array inputs
+        np_arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
+        mx_arr = mx.asarray(np_arr)
+        self.assertEqual(mx_arr.tolist(), [1.0, 2.0, 3.0])
+        self.assertEqual(mx_arr.dtype, mx.float32)
+
+        # dtype parameter
+        self.assertEqual(mx.asarray([1, 2, 3], dtype=mx.float32).dtype, mx.float32)
+        self.assertEqual(mx.asarray(42, dtype=mx.float16).dtype, mx.float16)
+
     def test_to_scalar(self):
         a = mx.array(1)
         self.assertEqual(int(a), 1)

python/tests/test_blas.py

@@ -297,6 +297,7 @@ class TestBlas(mlx_tests.MLXTestCase):
             self.assertTrue(np.allclose(out_mlx, out_npy, atol=1e-5))
 
     def test_matrix_vector(self):
+        mx.random.seed(0)
         for dtype in self.dtypes:
             with self.subTest(dtype=dtype):
                 np_dtype = getattr(np, dtype)

python/tests/test_load.py

@@ -168,8 +168,8 @@ class TestLoad(mlx_tests.MLXTestCase):
 
         expected = [
             0,
-            mx.nan,
-            mx.nan,
+            448,
+            -448,
             -0.875,
             0.4375,
             -0.005859,
@@ -179,12 +179,12 @@ class TestLoad(mlx_tests.MLXTestCase):
             -0.0039,
         ]
         expected = mx.array(expected, dtype=mx.bfloat16)
-        contents = b'H\x00\x00\x00\x00\x00\x00\x00{"tensor":{"dtype":"F8_E4M3","shape":[10],"data_offsets":[0,10]}}       \x00\x7f\xff\xb6.\x83\xba\xba\xbc\x82'
+        contents = b'H\x00\x00\x00\x00\x00\x00\x00{"tensor":{"dtype":"F8_E4M3","shape":[10],"data_offsets":[0,10]}}       \x00~\xfe\xb6.\x83\xba\xba\xbc\x82'
         with tempfile.NamedTemporaryFile(suffix=".safetensors") as f:
             f.write(contents)
             f.seek(0)
             out = mx.load(f)["tensor"]
-        self.assertTrue(mx.allclose(out[0], expected[0], equal_nan=True))
+        self.assertTrue(mx.allclose(mx.from_fp8(out), expected))
 
     def test_save_and_load_gguf_metadata_basic(self):
         if not os.path.isdir(self.test_dir):

python/tests/test_ops.py

@@ -3197,8 +3197,6 @@ class TestOps(mlx_tests.MLXTestCase):
             )
         )
 
-
-class TestBroadcast(mlx_tests.MLXTestCase):
     def test_broadcast_shapes(self):
         # Basic broadcasting
         self.assertEqual(mx.broadcast_shapes((1, 2, 3), (3,)), (1, 2, 3))
@@ -3243,6 +3241,13 @@ class TestBroadcast(mlx_tests.MLXTestCase):
         self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
         x = mx.array([3.0, mx.nan, 2.0, 0.0]) + 1j * mx.array([1.0] * 4)
 
+    def test_to_from_fp8(self):
+        vals = mx.array(
+            [448, 256, 192, 128, 96, 64, 48, 32, 24, 16, 12, 8, 6, 4, 3, 2, 0.015625]
+        )
+        self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals))
+        self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

tests/compile_tests.cpp

@@ -803,3 +803,16 @@ TEST_CASE("test compile with no-ops") {
   auto out = compile(fun)({in})[0];
   CHECK_EQ(out.inputs()[0].id(), in.id());
 }
+
+TEST_CASE("test compile random bits") {
+  auto fun = [](const std::vector<array>& inputs) {
+    auto key = inputs[0];
+    auto a = random::bits({32, 32}, 4, key);
+    auto b = random::bits({32, 32}, 2, key);
+    return std::vector<array>{a + b};
+  };
+  auto in = random::key(0);
+  auto expected = fun({in})[0];
+  auto out = compile(fun)({in})[0];
+  CHECK(array_equal(out, expected).item<bool>());
+}