From 1b081e2631517b3f2692f06ae9419a2d7dbb8b40 Mon Sep 17 00:00:00 2001
From: Huan <huan_xu@apple.com>
Date: Mon, 2 Mar 2026 23:27:27 -0800
Subject: [PATCH] Use num bytes, not num elems for buffer_size

---
 mlx/backend/metal/device.cpp | 2 +-
 mlx/backend/metal/device.h   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp
index 15824d6c5c..f4911a664e 100644
--- a/mlx/backend/metal/device.cpp
+++ b/mlx/backend/metal/device.cpp
@@ -259,7 +259,7 @@ void CommandEncoder::set_input_array(
     int idx,
     int64_t offset /* = 0 */) {
   if (all_inputs_.insert(a.buffer().ptr()).second) {
-    stream_.buffer_sizes += a.data_size();
+    stream_.buffer_sizes += a.data_size() * a.itemsize();
   }
   auto r_buf = static_cast<MTL::Resource*>(const_cast<void*>(a.buffer().ptr()));
   needs_barrier_ =
diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h
index e6162d7da6..fe5e89a3b5 100644
--- a/mlx/backend/metal/device.h
+++ b/mlx/backend/metal/device.h
@@ -137,7 +137,7 @@ struct DeviceStream {
   // Data updated between command buffers
   MTL::CommandBuffer* buffer{nullptr};
   int buffer_ops{0};
-  size_t buffer_sizes{0};
+  size_t buffer_sizes{0}; // bytes referenced by the command buffer
 
   // The command encoder, fence, and temporaries are updated between command
   // encoders