From 1e855446b25d6b45f9f54d9a4db29e605b888adb Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 17 Mar 2026 07:10:12 +0900 Subject: [PATCH] [CUDA] Pipelined QMM (#3255) --- mlx/backend/cuda/quantized/qmm/CMakeLists.txt | 3 + mlx/backend/cuda/quantized/qmm/qmm.cu | 75 +++ mlx/backend/cuda/quantized/qmm/qmm.h | 22 + .../cuda/quantized/qmm/qmm_impl_sm80.cuh | 482 ++++++++++++++++++ .../cuda/quantized/qmm/qmm_impl_sm80_m16.cu | 5 + .../cuda/quantized/qmm/qmm_impl_sm80_m32.cu | 5 + .../cuda/quantized/qmm/qmm_impl_sm80_m64.cu | 5 + mlx/backend/cuda/quantized/quantized.cpp | 60 ++- 8 files changed, 629 insertions(+), 28 deletions(-) create mode 100644 mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh create mode 100644 mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu create mode 100644 mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu create mode 100644 mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu diff --git a/mlx/backend/cuda/quantized/qmm/CMakeLists.txt b/mlx/backend/cuda/quantized/qmm/CMakeLists.txt index 35162a6b..9e057866 100644 --- a/mlx/backend/cuda/quantized/qmm/CMakeLists.txt +++ b/mlx/backend/cuda/quantized/qmm/CMakeLists.txt @@ -3,6 +3,9 @@ target_sources( PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu + ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m16.cu + ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m32.cu + ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m64.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n16_m1.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n32_m1.cu ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n64_m2.cu diff --git a/mlx/backend/cuda/quantized/qmm/qmm.cu b/mlx/backend/cuda/quantized/qmm/qmm.cu index 70b95fc7..6d5fa188 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm.cu +++ b/mlx/backend/cuda/quantized/qmm/qmm.cu @@ -97,6 +97,81 @@ void qmm_sm90( #endif // defined(MLX_CUDA_SM90A_ENABLED) } +// Defined in qmm_impl_sm80_xxx.cu files. +template +void qmm_impl_sm80( + const array& x, + const array& w, + const array& scales, + const array& biases, + array& out, + int bits, + int group_size, + cu::CommandEncoder& encoder); + +bool supports_qmm_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device) { + if (device.compute_capability_major() < 8) { + return false; + } + int n = out.shape(-1); + int k = x.shape(-1); + if ((n % 128 != 0) || (k % std::max(64, group_size) != 0)) { + return false; + } + if (!biases) { + return false; + } + if (!x.flags().row_contiguous || !w.flags().row_contiguous || + !scales.flags().row_contiguous || !biases->flags().row_contiguous) { + return false; + } + if (x.dtype() != float16 && x.dtype() != bfloat16) { + return false; + } + if (!transpose) { + return false; + } + if (bits != 8) { + return false; + } + if (mode != QuantizationMode::Affine) { + return false; + } + return true; +} + +void qmm_sm80( + const array& x, + const array& w, + const array& scales, + const array& biases, + array& out, + int bits, + int group_size, + cu::CommandEncoder& encoder) { + auto dispatch = [&]() { + qmm_impl_sm80(x, w, scales, biases, out, bits, group_size, encoder); + }; + int m = out.shape(-2); + if (m <= 16) { + dispatch.template operator()<16>(); + } else if (m <= 32) { + dispatch.template operator()<32>(); + } else { + dispatch.template operator()<64>(); + } +} + bool supports_fp_qmv( const array& x, const array& w, diff --git a/mlx/backend/cuda/quantized/qmm/qmm.h b/mlx/backend/cuda/quantized/qmm/qmm.h index efcd8eaf..f729abb8 100644 --- a/mlx/backend/cuda/quantized/qmm/qmm.h +++ b/mlx/backend/cuda/quantized/qmm/qmm.h @@ -32,6 +32,28 @@ void qmm_sm90( cu::CommandEncoder& encoder, Stream s); +bool supports_qmm_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::Device& device); + +void qmm_sm80( + const array& x, + const array& w, + const array& scales, + const array& biases, + array& out, + int bits, + int group_size, + cu::CommandEncoder& encoder); + bool supports_fp_qmv( const array& x, const array& w, diff --git a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh new file mode 100644 index 00000000..46bfe806 --- /dev/null +++ b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh @@ -0,0 +1,482 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/quantized/qmm/qmm.h" +#include "mlx/dtype_utils.h" + +#include +#include + +// clang-format off + +// We can't put kernel code in mlx::core due to name conflicts of "Shape". +namespace cutlass_gemm { + +using namespace cute; + +template +union SharedStorage { + struct { + ArrayEngine> A; + ArrayEngine> B; + } mainloop; + struct { + ArrayEngine> C; + } epilogue; +}; + +template +__device__ __forceinline__ void +dequant(const Q& w, const S& s, const Z& z, T out) { + // Scale must be one element. + CUTE_STATIC_ASSERT_V(cosize(s.layout()) == Int<1>{}); + CUTE_STATIC_ASSERT_V(cosize(z.layout()) == Int<1>{}); + // Quant must be contiguous. + auto layout = coalesce(w.layout()); + CUTE_STATIC_ASSERT_V(stride(layout) == Int<1>{}); + // Use cutlass for conversions. + constexpr int N = size(layout); + using Element = typename T::value_type; + using Quant = typename Q::value_type; + auto& w_vec = *(reinterpret_cast*>(raw_pointer_cast(w.data()))); + Element scale = s[0]; + Element zero_point = z[0]; + cutlass::NumericArrayConverter converter; + auto w_dq = converter(w_vec) * scale + zero_point; + copy(make_tensor(make_rmem_ptr(&w_dq), out.layout()), out); +} + +template +__global__ void qmm_sm80_kernel( + ProblemShape shape_MNKL, CtaTiler cta_tiler, + const Element* A, StrideA dA, SmemLayoutA sA_layout, TiledCopyA g2s_copy_a, S2RAtomA s2r_atom_a, + const Quant* B, StrideB dB, SmemLayoutB sB_layout, TiledCopyB g2s_copy_b, S2RAtomB s2r_atom_b, + Element* C, StrideC dC, SmemLayoutC sC_layout, TiledCopyC s2g_copy_c, R2SAtomC r2s_atom_c, + const Element* S, const Element* Z, LayoutS S_layout, G2RAtomS g2r_atom_s, TiledMma mma) { + CUTE_STATIC_ASSERT_V(size(g2s_copy_a) == size(mma)); + CUTE_STATIC_ASSERT_V(size(g2s_copy_b) == size(mma)); + CUTE_STATIC_ASSERT_V(size(s2g_copy_c) == size(mma)); + CUTE_STATIC_ASSERT_V(congruent(select<0,2,3>(shape_MNKL), dA)); + CUTE_STATIC_ASSERT_V(congruent(select<1,2,3>(shape_MNKL), dB)); + 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); + + // Represent the full tensors. + Tensor mA_mkl = make_tensor(make_gmem_ptr(A), select<0,2,3>(shape_MNKL), dA); // (M,K,L) + Tensor mB_nkl = make_tensor(make_gmem_ptr(B), select<1,2,3>(shape_MNKL), dB); // (N,K,L) + Tensor mC_mnl = make_tensor(make_gmem_ptr(C), select<0,1,3>(shape_MNKL), dC); // (M,N,L) + + Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) + Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) + + // Get batch slice. + Tensor mA = mA_mkl(_,_,l_coord); // (M,K) + Tensor mB = mB_nkl(_,_,l_coord); // (N,K) + Tensor mC = mC_mnl(_,_,l_coord); // (M,N) + + Tensor mS = mS_nkl(_,_,l_coord); // (N,(group_size,K/group_size)) + Tensor mZ = mZ_nkl(_,_,l_coord); // (N,(group_size,K/group_size)) + + // Get the appropriate blocks for this thread block. + auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k) + Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + + // Shared memory buffers. + extern __shared__ char shared_memory[]; + using SharedStorage = SharedStorage; + SharedStorage& smem = *reinterpret_cast(shared_memory); + Tensor sA = make_tensor(make_smem_ptr(smem.mainloop.A.begin()), sA_layout); // (BLK_M,BLK_K) + Tensor sB = make_tensor(make_smem_ptr(smem.mainloop.B.begin()), sB_layout); // (BLK_N,BLK_K) + Tensor sC = make_tensor(make_smem_ptr(smem.epilogue.C.begin()), sC_layout); // (BLK_M,BLK_N) + + // Partition the copying of A/B/C tiles across the threads. + ThrCopy g2s_thr_copy_a = g2s_copy_a.get_slice(thread_idx); + Tensor tAgA = g2s_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = g2s_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE) + + ThrCopy g2s_thr_copy_b = g2s_copy_b.get_slice(thread_idx); + Tensor tBgB = g2s_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = g2s_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE) + + ThrCopy s2g_thr_copy_c = s2g_copy_c.get_slice(thread_idx); + Tensor s2g_tCsC = s2g_thr_copy_c.partition_S(sC); // (CCPY,CCPY_M,CCPY_N) + Tensor s2g_tCgC = s2g_thr_copy_c.partition_D(gC); // (CCPY,CCPY_M,CCPY_N) + + // MMA. + ThrMMA thr_mma = mma.get_slice(thread_idx); + Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0)); // (MMA,MMA_M,MMA_K) + Tensor tCsB = thr_mma.partition_B(sB(_,_,0)); // (MMA,MMA_N,MMA_K) + Tensor tCrB = make_fragment_like(tCsB); // (MMA,MMA_N,MMA_K) + Tensor tCrB_dq = make_fragment_like(tCsB); // (MMA,MMA_N,MMA_K) + Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N) + Tensor tCrC_accu = make_fragment_like(tCgC); // (MMA,MMA_M,MMA_N) + Tensor tCrC = make_fragment_like(tCgC); // (MMA,MMA_M,MMA_N) + + Tensor tCgS = thr_mma.partition_B(gS); // (MMA,MMA_N,MMA_K,k) + Tensor tCrS = make_tensor_like(tCgS(_,_,_,0)); // (MMA,MMA_N,MMA_K) + Tensor tCgZ = thr_mma.partition_B(gZ); // (MMA,MMA_N,MMA_K,k) + Tensor tCrZ = make_tensor_like(tCgZ(_,_,_,0)); // (MMA,MMA_N,MMA_K) + + // Copy Atom retiling. + TiledCopy s2r_copy_a = make_tiled_copy_A(s2r_atom_a, mma); + ThrCopy s2r_thr_copy_a = s2r_copy_a.get_slice(thread_idx); + Tensor s2r_tCsA = s2r_thr_copy_a.partition_S(sA); // (ACPY,MMA_M,MMA_K,PIPE) + Tensor s2r_tCrA = s2r_thr_copy_a.retile_D(tCrA); // (ACPY,MMA_M,MMA_K) + + TiledCopy s2r_copy_b = make_tiled_copy_B(s2r_atom_b, mma); + ThrCopy s2r_thr_copy_b = s2r_copy_b.get_slice(thread_idx); + Tensor s2r_tCsB = s2r_thr_copy_b.partition_S(sB); // (BCPY,MMA_N,MMA_K,PIPE) + Tensor s2r_tCrB = s2r_thr_copy_b.retile_D(tCrB); // (BCPY,MMA_N,MMA_K) + + TiledCopy r2s_copy_c = make_tiled_copy_C(r2s_atom_c, mma); + ThrCopy r2s_thr_copy_c = r2s_copy_c.get_slice(thread_idx); + Tensor r2s_tCrC = r2s_thr_copy_c.retile_S(tCrC); // (CCPY,MMA_M,MMA_N) + Tensor r2s_tCsC = r2s_thr_copy_c.partition_D(sC); // (CCPY,MMA_M,MMA_N) + + TiledCopy g2r_copy_s = make_tiled_copy_B(g2r_atom_s, mma); + ThrCopy g2r_thr_copy_s = g2r_copy_s.get_slice(thread_idx); + Tensor g2r_tCgS = g2r_thr_copy_s.partition_S(gS); // (BCPY,MMA_N,MMA_K,k) + Tensor g2r_tCrS = g2r_thr_copy_s.retile_D(tCrS); // (BCPY,MMA_N,MMA_K) + Tensor g2r_tCgZ = g2r_thr_copy_s.partition_S(gZ); // (BCPY,MMA_N,MMA_K,k) + Tensor g2r_tCrZ = g2r_thr_copy_s.retile_D(tCrZ); // (BCPY,MMA_N,MMA_K) + + // Predicates for m bound. + auto m_max_coord = size<0>(shape_MNKL) - size<0>(gA) * m_coord; // M - BLK_M * m_coord + Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K) + Tensor tCpC = make_tensor(make_shape(size<1>(s2g_tCsC), size<2>(s2g_tCsC)), Stride<_1,_0>{}); // (CPY_M,CPY_N) + Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) + Tensor cC = make_identity_tensor(make_shape(size<0>(sC), size<1>(sC))); // (BLK_M,BLK_N) + Tensor tAcA = g2s_thr_copy_a.partition_D(cA); // (CPY,CPY_M,CPY_K) + Tensor tCcC = s2g_thr_copy_c.partition_D(cC); // (CPY,CPY_M,CPY_N) + CUTE_UNROLL + for (int m = 0; m < size<0>(tApA); ++m) { + tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord; + } + CUTE_UNROLL + for (int m = 0; m < size<0>(tCpC); ++m) { + tCpC(m,0) = get<0>(tCcC(0,m,0)) < m_max_coord; + } + + auto K_PIPE_MAX = size<3>(tAsA); + int smem_pipe_read = 0; + int smem_pipe_write = 0; + + // Copy A/B: GMEM => SMEM. + auto fetch_gmem = [&](int tile) { + copy_if(g2s_copy_a, tApA, tAgA(_,_,_,tile), tAsA(_,_,_,smem_pipe_write)); + copy(g2s_copy_b, tBgB(_,_,_,tile), tBsB(_,_,_,smem_pipe_write)); + cp_async_fence(); + smem_pipe_write = (smem_pipe_write + 1) % K_PIPE_MAX; + }; + // Copy S/Z: GMEM => RMEM. + auto fetch_scales = [&](int tile) { + copy(g2r_copy_s, g2r_tCgS(_,_,_,tile), g2r_tCrS); + copy(g2r_copy_s, g2r_tCgZ(_,_,_,tile), g2r_tCrZ); + }; + // Copy A/B: SMEM => RMEM. + auto fetch_smem = [&](auto block) { + copy(s2r_atom_a, s2r_tCsA(_,_,block,smem_pipe_read), s2r_tCrA(_,_,block)); + copy(s2r_atom_b, s2r_tCsB(_,_,block,smem_pipe_read), s2r_tCrB(_,_,block)); + CUTE_UNROLL + for (int n = 0; n < size<1>(tCrB); ++n) { + dequant(tCrB(_,n,block), tCrS(_,n,block), tCrZ(_,n,block), tCrB_dq(_,n,block)); + } + }; + + auto K_TILE_MAX = size<3>(tAgA); + auto K_BLOCK_MAX = size<2>(tCrA); + + // Prefetch beginning tiles. + CUTE_UNROLL + int tile_pipe = 0; + for (; tile_pipe < K_PIPE_MAX - 1; ++tile_pipe) { + fetch_gmem(tile_pipe); + } + + // Clear accumulators. + clear(tCrC_accu); + + // Prefetch first block. + if constexpr (K_BLOCK_MAX > 1) { + cp_async_wait(); + __syncthreads(); + fetch_scales(0); + fetch_smem(Int<0>{}); + } + + // Loop over CTA tiles. + for (int tile = 0; tile < K_TILE_MAX; ++tile) { + // Unroll MMA blocks. + CUTE_UNROLL + for (int block = 0; block < K_BLOCK_MAX; ++block) { + // Wait for last tile. + if (block == K_BLOCK_MAX - 1) { + smem_pipe_read = (smem_pipe_read + 1) % K_PIPE_MAX; + cp_async_wait(); + __syncthreads(); + fetch_scales((tile + 1 < K_TILE_MAX) ? tile + 1 : tile); + } + // Prefetch next block. + fetch_smem((block + 1) % K_BLOCK_MAX); + // Prefetch next tile. + if (block == 0) { + fetch_gmem(tile_pipe); + tile_pipe = (tile_pipe + 1 < K_TILE_MAX) ? tile_pipe + 1 : tile_pipe; + } + // MMA. + gemm(mma, tCrA(_,_,block), tCrB_dq(_,_,block), tCrC_accu); + } + } + + // Epilogue. + CUTE_UNROLL + for (int i = 0; i < size(tCrC_accu); i++) { + tCrC(i) = Element(tCrC_accu(i)); + } + copy(r2s_copy_c, r2s_tCrC, r2s_tCsC); + __syncthreads(); + copy_if(s2g_copy_c, tCpC, s2g_tCsC, s2g_tCgC); +} + +template +inline constexpr auto make_mma_atom() { + if constexpr (std::is_same_v) { + return SM80_16x8x16_F32F16F16F32_TN{}; + } + if constexpr (std::is_same_v) { + return SM80_16x8x16_F32BF16BF16F32_TN{}; + } +} + +template +inline constexpr auto make_tiled_mma() { + constexpr auto atom = make_mma_atom(); + if constexpr (TileM >= 32) { + return make_tiled_mma(atom, Layout>{}, Tile<_32,_32,_16>{}); + } else { + return make_tiled_mma(atom, Layout>{}, Tile<_16,_32,_16>{}); + } +} + +template typename Atom, typename NumThreads> +inline auto make_tiled_copy(NumThreads num_threads) { + return make_tiled_copy( + Copy_Atom>, T>{}, + make_layout(make_shape(Int{}, Int<8>{}), LayoutRight{}), + make_layout(make_shape(Int<1>{}, Int>{}))); +} + +template +void qmm_sm80( + const Element* A, + const Quant* B, + const Element* S, + const Element* Z, + Element* C, + int m, int n, int k, int l, + GroupSize group_size, + F&& launch_kernel) { + // Define shapes (dynamic). + auto prob_shape = make_shape(m, n, k, l); // (M,N,K,L) + + // Define TN strides (mixed). + auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL) + auto dB = make_stride(k, Int<1>{}, n * k); // (dN,dK,dL) + auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) + + // Define CTA tile sizes (static). + auto bM = Int{}; + auto bN = Int<128>{}; + auto bK = Int{}; + auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M,BLK_N,BLK_K) + + // Define MMA. + TiledMMA mma = make_tiled_mma(); + auto num_threads = size(mma); + + // Define the A/B smem layouts (static). + auto swizzle_ab = composition(Swizzle<3,3,3>{}, + Layout>, + Stride<_8,Stride<_1,_64>>>{}); + auto bP = Int<3>{}; // pipeline + auto sA_layout = tile_to_shape(swizzle_ab, make_shape(bM, bK, bP)); + auto sB_layout = tile_to_shape(swizzle_ab, make_shape(bN, bK, bP)); + + // Define the C smem layouts (static). + // TODO: Find a better swizzle. + auto sC_layout = tile_to_shape(swizzle_ab, make_shape(bM, bN)); + + // Define the scales/biases smem layouts (static). + auto bS = ceil_div(bK, group_size); + auto sS_layout = make_layout(make_shape(bN, make_shape(group_size, bS)), + make_stride(bS, Stride<_0, _1>{})); + + // Define layout of scales/biases (mixed). + auto S_layout = make_layout( + make_shape(n, make_shape(group_size, k / group_size), l), + make_stride(k / group_size, Stride<_0, _1>{}, n * k / group_size)); + + // Atoms. + constexpr int element_bits = sizeof_bits_v; + constexpr int quant_bits = sizeof_bits_v; + constexpr int qload = 128 / (element_bits / quant_bits); + TiledCopy g2s_copy_a = make_tiled_copy(num_threads); + TiledCopy g2s_copy_b = make_tiled_copy(num_threads); + TiledCopy s2g_copy_c = make_tiled_copy(num_threads); + + Copy_Atom s2r_atom_a; + Copy_Atom>, Quant> s2r_atom_b; + Copy_Atom>, Element> r2s_atom_c; + Copy_Atom, Element> g2r_atom_s; + + auto* kernel = &qmm_sm80_kernel< + decltype(prob_shape), decltype(cta_tiler), + Element, Quant, + decltype(dA), decltype(sA_layout), decltype(g2s_copy_a), decltype(s2r_atom_a), + decltype(dB), decltype(sB_layout), decltype(g2s_copy_b), decltype(s2r_atom_b), + decltype(dC), decltype(sC_layout), decltype(s2g_copy_c), decltype(r2s_atom_c), + decltype(S_layout), decltype(g2r_atom_s), decltype(mma)>; + + // Set L1 to be SMEM only. + size_t smem_bytes = sizeof(SharedStorage); + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100); + + dim3 num_blocks(size(ceil_div(m, bM)), size(ceil_div(n, bN)), l); + dim3 block_dims(num_threads); + void* args[] = { + &prob_shape, &cta_tiler, + &A, &dA, &sA_layout, &g2s_copy_a, &s2r_atom_a, + &B, &dB, &sB_layout, &g2s_copy_b, &s2r_atom_b, + &C, &dC, &sC_layout, &s2g_copy_c, &r2s_atom_c, + &S, &Z, &S_layout, &g2r_atom_s, &mma}; + launch_kernel(reinterpret_cast(kernel), num_blocks, block_dims, smem_bytes, args); +} + +} // namespace cutlass_gemm + +// clang-format on + +namespace mlx::core { + +template +inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) { + if (dtype == float16) { + f.template operator()(); + } else if (dtype == bfloat16) { + f.template operator()(); + } else { + throw std::invalid_argument( + fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype))); + } +} + +template +inline void dispatch_quant_types(int bits, const char* tag, F&& f) { + if (bits == 4) { + f.template operator()(); + } else if (bits == 8) { + f.template operator()(); + } else { + throw std::invalid_argument( + fmt::format("{} {}-bit quantization is not supported.", tag, bits)); + } +} + +template +inline void dispatch_groups(int group_size, const char* tag, F&& f) { + if (group_size == 32) { + f.template operator()<32>(); + } else if (group_size == 64) { + f.template operator()<64>(); + } else if (group_size == 128) { + f.template operator()<128>(); + } else { + throw std::invalid_argument( + fmt::format("{} Group size {} is not supported.", tag, group_size)); + } +} + +template +void qmm_impl_sm80( + const array& x, + const array& w, + const array& scales, + const array& biases, + array& out, + int bits, + int group_size, + cu::CommandEncoder& encoder) { + const char* tag = "[quantized_matmul]"; + int m = out.shape(-2); + int n = out.shape(-1); + int k = x.shape(-1); + int l = out.size() / (m * n); + + dispatch_element_types(out.dtype(), tag, [&]() { + dispatch_quant_types(bits, tag, [&]() { + dispatch_groups(group_size, tag, [&]() { + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + encoder.set_input_array(biases); + encoder.set_output_array(out); + cutlass_gemm::qmm_sm80( + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(scales), + gpu_ptr(biases), + gpu_ptr(out), + m, + n, + k, + l, + cute::Int{}, + [&](auto* kernel, + dim3 num_blocks, + dim3 block_dims, + uint32_t smem_bytes, + void** args) { + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, smem_bytes, args); + }); + }); + }); + }); +} + +} // namespace mlx::core + +#define QMM_SM80_GPU(TileM) \ + namespace mlx::core { \ + template void qmm_impl_sm80( \ + const array& x, \ + const array& w, \ + const array& scales, \ + const array& biases, \ + array& out, \ + int bits, \ + int group_size, \ + cu::CommandEncoder& encoder); \ + } diff --git a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu new file mode 100644 index 00000000..cd682be3 --- /dev/null +++ b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu @@ -0,0 +1,5 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh" + +QMM_SM80_GPU(16) diff --git a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu new file mode 100644 index 00000000..1f79364b --- /dev/null +++ b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu @@ -0,0 +1,5 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh" + +QMM_SM80_GPU(32) diff --git a/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu new file mode 100644 index 00000000..9af44f8e --- /dev/null +++ b/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu @@ -0,0 +1,5 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh" + +QMM_SM80_GPU(64) diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp index 62a067bf..03820511 100644 --- a/mlx/backend/cuda/quantized/quantized.cpp +++ b/mlx/backend/cuda/quantized/quantized.cpp @@ -25,19 +25,6 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { biases = inputs[3]; } - 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); - }; - auto call_fp_qmv = [&]() { - out.set_data(cu::malloc_async(out.nbytes(), encoder)); - fp_qmv(x, w, scales, out, bits_, group_size_, encoder, s); - }; - auto call_qmv = [&]() { - out.set_data(cu::malloc_async(out.nbytes(), encoder)); - qmv(x, w, scales, biases, out, bits_, group_size_, mode_, encoder); - }; - auto supports = [&](auto&& f) { return f( x, @@ -52,34 +39,50 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { encoder.device()); }; bool can_use_qmm_sm90 = supports(supports_qmm_sm90); + bool can_use_qmm_sm80 = supports(supports_qmm_sm80); bool can_use_fp_qmv = supports(supports_fp_qmv); - bool can_use_qmv = supports(supports_qmv); + bool can_use_qmv = supports(supports_qmv) || can_use_fp_qmv; + + 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); + }; + auto call_qmm_sm80 = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + qmm_sm80(x, w, scales, *biases, out, bits_, group_size_, 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, w, scales, biases, out, bits_, group_size_, mode_, encoder); + } + }; int M = out.shape(-2); int N = out.shape(-1); int K = x.shape(-1); int B = out.size() / (M * N); - bool prefer_qmv = M == 1 && B == 1 && N <= 16384 && K <= 16384; if (can_use_qmm_sm90) { - if (prefer_qmv) { - if (can_use_fp_qmv) { - call_fp_qmv(); - return; - } - if (can_use_qmv) { - call_qmv(); - return; - } + if (can_use_qmv && (M == 1 && B == 1 && N <= 16384 && K <= 16384)) { + call_qmv(); + } else { + call_qmm_sm90(); } - call_qmm_sm90(); return; } - if (can_use_fp_qmv) { - call_fp_qmv(); + if (can_use_qmm_sm80) { + if (can_use_qmv && (M * B < 8)) { + call_qmv(); + } else { + call_qmm_sm80(); + } return; } + if (can_use_qmv) { call_qmv(); return; @@ -88,12 +91,13 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { throw std::runtime_error( fmt::format( "[quantized_matmul] No implementation for " - "problem shape: {}x{}x{}x{} " + "problem shape: {}x{}x{}x{}, transpose: {}, " "activation: {}, bits: {}, group size: {}, mode: \"{}\".", M, N, K, B, + transpose_, dtype_to_string(x.dtype()), bits_, group_size_,