diff --git a/mlx/backend/cuda/random.cu b/mlx/backend/cuda/random.cu
index 72f6c84222..1a485134d9 100644
--- a/mlx/backend/cuda/random.cu
+++ b/mlx/backend/cuda/random.cu
@@ -49,7 +49,7 @@ __global__ void rbitsc(
     uint8_t* out,
     dim3 grid_dims,
     bool odd,
-    uint32_t bytes_per_key) {
+    uint64_t bytes_per_key) {
   auto grid = cg::this_grid();
   uint32_t thread_index = grid.thread_rank();
   uint32_t index_x = thread_index % grid_dims.x;
@@ -89,7 +89,7 @@ __global__ void rbits(
     uint8_t* out,
     dim3 grid_dims,
     bool odd,
-    uint32_t bytes_per_key,
+    uint64_t bytes_per_key,
     int32_t ndim,
     const __grid_constant__ Shape key_shape,
     const __grid_constant__ Strides key_strides) {
diff --git a/mlx/backend/metal/kernels/random.metal b/mlx/backend/metal/kernels/random.metal
index ccbd464d3b..13d7fd931a 100644
--- a/mlx/backend/metal/kernels/random.metal
+++ b/mlx/backend/metal/kernels/random.metal
@@ -35,7 +35,7 @@ rbits threefry2x32_hash(const thread uint2& key, uint2 count) {
     device const uint32_t* keys,
     device char* out,
     constant const bool& odd,
-    constant const uint& bytes_per_key,
+    constant const ulong& bytes_per_key,
     uint2 grid_dim [[threads_per_grid]],
     uint2 index [[thread_position_in_grid]]) {
   auto kidx = 2 * index.x;
@@ -68,7 +68,7 @@ rbits threefry2x32_hash(const thread uint2& key, uint2 count) {
     device const uint32_t* keys,
     device char* out,
     constant const bool& odd,
-    constant const uint& bytes_per_key,
+    constant const ulong& bytes_per_key,
     constant const int& ndim,
     constant const int* key_shape,
     constant const int64_t* key_strides,