From 80bcd1c65867dec1623546d1fdc9d0473ab46657 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 6 May 2026 08:35:53 +0900 Subject: [PATCH] [CUDA] Fix half type matmul in cutlass kernels (#3469) --- mlx/backend/cuda/gemms/gather_gemm.cu | 15 ++++++++++----- .../cuda/gemms/grouped_gemm_unaligned.cu | 2 +- mlx/backend/cuda/quantized/qmm/qmm.cu | 8 ++++---- mlx/backend/cuda/quantized/qmm/qmm_naive.cu | 4 +++- mlx/backend/cuda/quantized/qmm/qmm_sm80.cu | 4 +++- mlx/backend/cuda/quantized/quantized.cpp | 10 ++++++---- python/tests/cuda_skip.py | 17 ----------------- python/tests/test_quantized.py | 2 +- 8 files changed, 28 insertions(+), 34 deletions(-) diff --git a/mlx/backend/cuda/gemms/gather_gemm.cu b/mlx/backend/cuda/gemms/gather_gemm.cu index 2c971be4..fbbbf159 100644 --- a/mlx/backend/cuda/gemms/gather_gemm.cu +++ b/mlx/backend/cuda/gemms/gather_gemm.cu @@ -43,6 +43,12 @@ class GatherGemm { using ElementD = typename CollectiveEpilogue::ElementD; using StrideD = typename CollectiveEpilogue::StrideD; + static_assert( + cute::is_same_v< + ElementAccumulator, + typename CollectiveEpilogue::ElementAccumulator>, + "Mainloop and epilogue do not agree on accumulator value type."); + static constexpr int SharedStorageSize = static_cast(cute::max( sizeof(typename CollectiveMainloop::SharedStorage), sizeof(typename CollectiveEpilogue::SharedStorage))); @@ -98,7 +104,9 @@ class GatherGemm { CUTLASS_DEVICE void operator()(const Params& params, char* smem_buf) { int thread_idx = int(threadIdx.x); - auto [m_coord, n_coord, l_coord] = uint3(blockIdx); + int m_coord = int(blockIdx.x); + int n_coord = int(blockIdx.y); + int l_coord = int(blockIdx.z); auto shape_MNKL = append<4>(params.problem_shape, Int<1>{}); auto cta_tile = TileShape{}; @@ -220,7 +228,7 @@ void gather_mm( using TileShape = Shape<_128, _128, _8>; using DispatchPolicy = cutlass::gemm::MainloopSm70TwoStage; using TiledMma = TiledMMA< - MMA_Atom>, + MMA_Atom>, Layout>>; using CopyTraitsA = SimtCopyTraits; @@ -296,9 +304,6 @@ void cutlass_gather_mm( int n = out.shape(-1); int k = a.shape(-1); int l = out.size() / (m * n); - if (m < 16 || n < 16) { - throw std::invalid_argument("[gather_mm] M/N is too small."); - } encoder.set_input_array(a); encoder.set_input_array(b); diff --git a/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu b/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu index 08b3f1c1..319c759c 100644 --- a/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu +++ b/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu @@ -245,7 +245,7 @@ void grouped_gemm_v2( LayoutB, cutlass::ComplexTransform::kNone, GemmConfiguration::kAlignmentAB, - typename GemmConfiguration::Element, + typename GemmConfiguration::Accumulator, cutlass::layout::RowMajor, typename GemmConfiguration::Accumulator, typename GemmConfiguration::OpClass, diff --git a/mlx/backend/cuda/quantized/qmm/qmm.cu b/mlx/backend/cuda/quantized/qmm/qmm.cu index 8385cbd5..41e802d6 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm.cu +++ b/mlx/backend/cuda/quantized/qmm/qmm.cu @@ -52,7 +52,7 @@ bool supports_qmm_sm90( if (!biases) { return false; } - if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) || + if (!is_last_2_dims_row_contiguous(w) || !is_last_2_dims_row_contiguous(scales) || !is_last_2_dims_row_contiguous(*biases)) { return false; @@ -139,7 +139,7 @@ bool supports_qmm_sm80( if ((n % 128 != 0) || (k % std::max(64, group_size) != 0)) { return false; } - if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) || + if (!is_last_2_dims_row_contiguous(w) || !is_last_2_dims_row_contiguous(scales)) { return false; } @@ -224,7 +224,7 @@ bool supports_qmm_naive( if (transpose && (k % std::max(64, group_size) != 0)) { return false; } - if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) || + if (!is_last_2_dims_row_contiguous(w) || !is_last_2_dims_row_contiguous(scales)) { return false; } @@ -343,7 +343,7 @@ bool supports_qmv( if (k % 8 != 0) { return false; } - if (!x.flags().row_contiguous || !is_last_2_dims_row_contiguous(w) || + if (!is_last_2_dims_row_contiguous(w) || !is_last_2_dims_row_contiguous(scales)) { return false; } diff --git a/mlx/backend/cuda/quantized/qmm/qmm_naive.cu b/mlx/backend/cuda/quantized/qmm/qmm_naive.cu index 5be75bd0..5e23dab4 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm_naive.cu +++ b/mlx/backend/cuda/quantized/qmm/qmm_naive.cu @@ -60,7 +60,9 @@ __global__ void qmm_naive_kernel( CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC)); int thread_idx = int(threadIdx.x); - auto [m_coord, n_coord, l_coord] = static_cast(blockIdx); + int m_coord = int(blockIdx.x); + int n_coord = int(blockIdx.y); + int l_coord = int(blockIdx.z); auto m_max_coord = size<0>(shape_MNKL) - size<0>(cta_tiler) * m_coord; // M - BLK_M * m_coord auto n_max_coord = size<1>(shape_MNKL) - size<1>(cta_tiler) * n_coord; // N - BLK_N * n_coord diff --git a/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu b/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu index 028f18cd..b721c0d7 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu +++ b/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu @@ -48,7 +48,9 @@ __global__ void qmm_sm80_kernel( CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC)); int thread_idx = int(threadIdx.x); - auto [m_coord, n_coord, l_coord] = static_cast(blockIdx); + int m_coord = int(blockIdx.x); + int n_coord = int(blockIdx.y); + int l_coord = int(blockIdx.z); // For gather, use index lookup for input batch slicing. uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord; diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index c3ac09b1..645b24cc 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/mlx/backend/cuda/quantized/quantized.cpp @@ -17,7 +17,7 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& encoder = cu::get_command_encoder(s); - const array& x = inputs[0]; + array x = ensure_row_contiguous(inputs[0], encoder, s); const array& w = inputs[1]; const array& scales = inputs[2]; std::optional biases; @@ -146,15 +146,17 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& encoder = cu::get_command_encoder(s); - const array& x = inputs[0]; + array x = ensure_row_contiguous(inputs[0], encoder, s); const array& w = inputs[1]; const array& scales = inputs[2]; std::optional biases; if (inputs.size() == 6) { biases = inputs[3]; } - array lhs_indices = ensure_contiguous(inputs[inputs.size() - 2], encoder, s); - array rhs_indices = ensure_contiguous(inputs[inputs.size() - 1], encoder, s); + array lhs_indices = + ensure_row_contiguous(inputs[inputs.size() - 2], encoder, s); + array rhs_indices = + ensure_row_contiguous(inputs[inputs.size() - 1], encoder, s); int M = out.ndim() > 1 ? out.shape(-2) : 1; int N = out.shape(-1); diff --git a/python/tests/cuda_skip.py b/python/tests/cuda_skip.py index 31a657d7..7f9f75f1 100644 --- a/python/tests/cuda_skip.py +++ b/python/tests/cuda_skip.py @@ -1,22 +1,5 @@ cuda_skip = { - # Lapack ops NYI - "TestLinalg.test_cholesky", - "TestLinalg.test_cholesky_inv", - "TestLinalg.test_eig", - "TestLinalg.test_eigh", - "TestLinalg.test_inverse", - "TestVmap.test_vmap_inverse", - "TestLinalg.test_lu", - "TestLinalg.test_lu_factor", - "TestLinalg.test_pseudo_inverse", - "TestLinalg.test_qr_factorization", - "TestInit.test_orthogonal", - "TestLinalg.test_svd_decomposition", - "TestVmap.test_vmap_svd", - "TestLinalg.test_tri_inverse", # Quantization NYI "TestQuantized.test_gather_matmul_grad", "TestQuantized.test_gather_qmm", - "TestQuantized.test_gather_qmm_sorted", - "TestQuantized.test_gather_qmm_grad", } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index a7472e99..afbd14bf 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1046,7 +1046,7 @@ class TestQuantized(mlx_tests.MLXTestCase): y3 = scatter_unsort(y3, inv_order, indices.shape) y4 = scatter_unsort(y4, inv_order, indices.shape) - tol = 1.5e-5 if (dtype == mx.float32) else 2.5e-4 + tol = 1.5e-5 if (dtype == mx.float32) else 1e-3 self.assertLess((y1 - y2).abs().max(), tol) self.assertLess((y1 - y3).abs().max(), tol)