fix(cuda/quantized): chunk long-prompt qmm launches to avoid gridDim>65535 and int overflow

src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp

@@ -1,10 +1,12 @@
 // Copyright © 2025 Apple Inc.
-// Patched by mlxcel: ensure input contiguity in QuantizedMatmul for
+// Patched by mlxcel: (1) ensure input contiguity in QuantizedMatmul for
 // non-contiguous 3D batched weights (e.g. GLM-4 MLA embed_q with
-// transpose=false). Synced to upstream e9463bb (post-#3706/#3576 JIT
-// qmm rework and #3723 qmv global scale; the dispatch consumed here kept
-// its public signatures, with qmv gaining an optional global_scale that
-// QuantizedMatmul passes as std::nullopt).
+// transpose=false); (2) split long-prompt quantized matmuls so no CUDA launch
+// exceeds the gridDim.y/z limit of 65535 and no `l = out.size()/(m*n)` int32
+// multiply overflows (see lablup/mlxcel#648). Synced to upstream e9463bb
+// (post-#3706/#3576 JIT qmm rework and #3723 qmv global scale; the dispatch
+// consumed here kept its public signatures, with qmv gaining an optional
+// global_scale that QuantizedMatmul passes as std::nullopt).
 
 #include "mlx/backend/cuda/quantized/quantized.h"
 #include "mlx/backend/cuda/device.h"
@@ -16,8 +18,99 @@
 
 #include <nvtx3/nvtx3.hpp>
 
+#include <algorithm>
+#include <cstdint>
+
 namespace mlx::core {
 
+namespace {
+
+// CUDA gridDim.y and gridDim.z are capped at 65535. The quantized-matmul
+// kernels place the row count (m) in grid.y and the batch/gather count (l) in
+// grid.z, and qmv/make_problem_shape derive `l = out.size() / (m * n)` with m,n
+// as int. Two long-prompt failure modes result (lablup/mlxcel#648): (a) MoE
+// GatherQMM makes l = tokens * num_experts_per_tok, which exceeds 65535 once
+// tokens*top_k >= 65536; (b) a dense LM-head qmv has m*n = tokens*vocab, which
+// overflows int32 once tokens*vocab >= 2^31, so l wraps to 0 and grid.z = 0 is
+// rejected as an invalid launch. Both are avoided by splitting the leading
+// (row/batch) dimension into slices small enough that every launch keeps its
+// grid dims <= 65535 and its m*n < 2^31.
+constexpr int64_t kMaxGridDim = 65535;
+
+// [count, inner] row-contiguous view of `src` at leading element offset
+// start*inner, sharing src's device buffer (no copy).
+array row_view(const array& src, int64_t start, int64_t count, int64_t inner) {
+  array v(Shape{static_cast<int>(count), static_cast<int>(inner)}, src.dtype(),
+          nullptr, {});
+  v.copy_shared_buffer(
+      src, Strides{inner, 1}, src.flags(), count * inner, start * inner);
+  return v;
+}
+
+// [bc, M, N] row-contiguous view selecting flat batch [b0, b0+bc) of `src`.
+array batch_view(
+    const array& src, int64_t b0, int64_t bc, int64_t M, int64_t N) {
+  array v(Shape{static_cast<int>(bc), static_cast<int>(M), static_cast<int>(N)},
+          src.dtype(), nullptr, {});
+  v.copy_shared_buffer(
+      src, Strides{M * N, N, 1}, src.flags(), bc * M * N, b0 * M * N);
+  return v;
+}
+
+// [bc] flat view selecting elements [b0, b0+bc) of `src`.
+array flat_view(const array& src, int64_t b0, int64_t bc) {
+  array v(Shape{static_cast<int>(bc)}, src.dtype(), nullptr, {});
+  v.copy_shared_buffer(src, Strides{1}, src.flags(), bc, b0);
+  return v;
+}
+
+// Invoke call_one(x, out) split into row chunks so grid.y stays <= 65535 and
+// each chunk's m*n stays < 2^31. Only the unbatched 2D path (the case that
+// overflows) is chunked; batched weights keep M small and pass through.
+template <typename F>
+void run_row_chunked(
+    const array& x, array& out, int M, int N, int K, F&& call_one) {
+  int64_t cap = std::min<int64_t>(
+      kMaxGridDim, static_cast<int64_t>(INT32_MAX) / std::max(N, 1));
+  // Only the single-batch (B == 1) case is split along its M (row) axis: the
+  // buffer is then exactly M*N contiguous regardless of any leading size-1 dims
+  // (e.g. an LM head with out shape [1, tokens, vocab]). Batched weights (B > 1)
+  // keep M small, cannot overflow, and pass through unchanged.
+  bool single_batch = static_cast<size_t>(out.size()) ==
+      static_cast<size_t>(M) * static_cast<size_t>(N);
+  if (!single_batch || static_cast<int64_t>(M) <= cap) {
+    call_one(x, out);
+    return;
+  }
+  for (int64_t r0 = 0; r0 < M; r0 += cap) {
+    int64_t rc = std::min<int64_t>(cap, static_cast<int64_t>(M) - r0);
+    array xc = row_view(x, r0, rc, K);
+    array oc = row_view(out, r0, rc, N);
+    call_one(xc, oc);
+  }
+}
+
+// Invoke call_one(lhs, rhs, out) split into batch chunks so grid.z (== l) stays
+// <= 65535. lhs/rhs indices are treated as flat [B] arrays.
+template <typename F>
+void run_batch_chunked(
+    array& out, const array& lhs, const array& rhs, int M, int N, F&& call_one) {
+  int64_t B = static_cast<int64_t>(out.size()) / M / N;
+  if (B <= kMaxGridDim) {
+    call_one(lhs, rhs, out);
+    return;
+  }
+  for (int64_t b0 = 0; b0 < B; b0 += kMaxGridDim) {
+    int64_t bc = std::min<int64_t>(kMaxGridDim, B - b0);
+    array oc = batch_view(out, b0, bc, M, N);
+    array lc = flat_view(lhs, b0, bc);
+    array rc = flat_view(rhs, b0, bc);
+    call_one(lc, rc, oc);
+  }
+}
+
+} // namespace
+
 void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   nvtx3::scoped_range r("QuantizedMatmul::eval_gpu");
   auto& s = stream();
@@ -54,63 +147,71 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   bool can_use_fp_qmv = supports(supports_fp_qmv);
   bool can_use_qmv = supports(supports_qmv) || can_use_fp_qmv;
 
+  int M = out.ndim() > 1 ? out.shape(-2) : 1;
+  int N = out.shape(-1);
+  int K = x.shape(-1);
+  int B = out.size() / M / N;
+
   auto call_qmm_sm90 = [&]() {
     out.set_data(cu::malloc_async(out.nbytes(), encoder));
-    qmm_sm90(x, w, scales, *biases, out, bits_, group_size_, encoder, s);
+    run_row_chunked(x, out, M, N, K, [&](const array& xc, array& oc) {
+      qmm_sm90(xc, w, scales, *biases, oc, bits_, group_size_, encoder, s);
+    });
   };
   auto call_qmm_sm80 = [&]() {
     out.set_data(cu::malloc_async(out.nbytes(), encoder));
-    qmm_sm80(
-        x,
-        w,
-        scales,
-        biases,
-        std::nullopt,
-        std::nullopt,
-        out,
-        bits_,
-        group_size_,
-        mode_,
-        encoder);
+    run_row_chunked(x, out, M, N, K, [&](const array& xc, array& oc) {
+      qmm_sm80(
+          xc,
+          w,
+          scales,
+          biases,
+          std::nullopt,
+          std::nullopt,
+          oc,
+          bits_,
+          group_size_,
+          mode_,
+          encoder);
+    });
   };
   auto call_qmm_naive = [&]() {
     out.set_data(cu::malloc_async(out.nbytes(), encoder));
-    qmm_naive(
-        x,
-        w,
-        scales,
-        biases,
-        std::nullopt,
-        std::nullopt,
-        out,
-        transpose_,
-        bits_,
-        group_size_,
-        mode_,
-        encoder);
-  };
-  auto call_qmv = [&]() {
-    out.set_data(cu::malloc_async(out.nbytes(), encoder));
-    if (can_use_fp_qmv) {
-      fp_qmv(x, w, scales, out, bits_, group_size_, encoder, s);
-    } else {
-      qmv(x,
+    run_row_chunked(x, out, M, N, K, [&](const array& xc, array& oc) {
+      qmm_naive(
+          xc,
           w,
           scales,
           biases,
           std::nullopt,
-          out,
+          std::nullopt,
+          oc,
+          transpose_,
           bits_,
           group_size_,
           mode_,
           encoder);
-    }
+    });
+  };
+  auto call_qmv = [&]() {
+    out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    run_row_chunked(x, out, M, N, K, [&](const array& xc, array& oc) {
+      if (can_use_fp_qmv) {
+        fp_qmv(xc, w, scales, oc, bits_, group_size_, encoder, s);
+      } else {
+        qmv(xc,
+            w,
+            scales,
+            biases,
+            std::nullopt,
+            oc,
+            bits_,
+            group_size_,
+            mode_,
+            encoder);
+      }
+    });
   };
-
-  int M = out.ndim() > 1 ? out.shape(-2) : 1;
-  int N = out.shape(-1);
-  int K = x.shape(-1);
-  int B = out.size() / (M * N);
 
   if (can_use_qmm_sm90) {
     if (can_use_qmv && (M == 1 && B == 1 && N <= 16384 && K <= 16384)) {
@@ -180,7 +281,7 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
   int M = out.ndim() > 1 ? out.shape(-2) : 1;
   int N = out.shape(-1);
   int K = x.shape(-1);
-  int B = out.size() / (M * N);
+  int B = out.size() / M / N;
 
   auto supports = [&](auto&& f) {
     return f(
@@ -201,49 +302,61 @@ void GatherQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
 
   auto call_qmm_sm80 = [&]() {
     out.set_data(cu::malloc_async(out.nbytes(), encoder));
-    qmm_sm80(
-        x,
-        w,
-        scales,
-        biases,
-        lhs_indices,
-        rhs_indices,
-        out,
-        bits_,
-        group_size_,
-        mode_,
-        encoder);
+    run_batch_chunked(
+        out, lhs_indices, rhs_indices, M, N,
+        [&](const array& lc, const array& rc, array& oc) {
+          qmm_sm80(
+              x,
+              w,
+              scales,
+              biases,
+              lc,
+              rc,
+              oc,
+              bits_,
+              group_size_,
+              mode_,
+              encoder);
+        });
   };
   auto call_qmm_naive = [&]() {
     out.set_data(cu::malloc_async(out.nbytes(), encoder));
-    qmm_naive(
-        x,
-        w,
-        scales,
-        biases,
-        lhs_indices,
-        rhs_indices,
-        out,
-        transpose_,
-        bits_,
-        group_size_,
-        mode_,
-        encoder);
+    run_batch_chunked(
+        out, lhs_indices, rhs_indices, M, N,
+        [&](const array& lc, const array& rc, array& oc) {
+          qmm_naive(
+              x,
+              w,
+              scales,
+              biases,
+              lc,
+              rc,
+              oc,
+              transpose_,
+              bits_,
+              group_size_,
+              mode_,
+              encoder);
+        });
   };
   auto call_qmv = [&]() {
     out.set_data(cu::malloc_async(out.nbytes(), encoder));
-    gather_qmv(
-        x,
-        w,
-        scales,
-        biases,
-        lhs_indices,
-        rhs_indices,
-        out,
-        bits_,
-        group_size_,
-        mode_,
-        encoder);
+    run_batch_chunked(
+        out, lhs_indices, rhs_indices, M, N,
+        [&](const array& lc, const array& rc, array& oc) {
+          gather_qmv(
+              x,
+              w,
+              scales,
+              biases,
+              lc,
+              rc,
+              oc,
+              bits_,
+              group_size_,
+              mode_,
+              encoder);
+        });
   };
 
   if (can_use_qmm_sm80) {