From be872ebdefbc769045e1fcde3e34b9e664488e13 Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Thu, 5 Mar 2026 16:34:19 +0800 Subject: [PATCH] [CUDA] implement Hadamard transform (#3179) --- mlx/backend/cuda/CMakeLists.txt | 1 + mlx/backend/cuda/device/hadamard.cuh | 184 ++++++++++++++++++++ mlx/backend/cuda/hadamard.cu | 245 +++++++++++++++++++++++++++ mlx/backend/cuda/jit_module.cpp | 2 + mlx/backend/cuda/primitives.cpp | 1 - python/tests/cuda_skip.py | 3 - 6 files changed, 432 insertions(+), 4 deletions(-) create mode 100644 mlx/backend/cuda/device/hadamard.cuh create mode 100644 mlx/backend/cuda/hadamard.cu diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index bc1bfa9f..e968b1c1 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -30,6 +30,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu + ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cu ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu diff --git a/mlx/backend/cuda/device/hadamard.cuh b/mlx/backend/cuda/device/hadamard.cuh new file mode 100644 index 00000000..9d2081e0 --- /dev/null +++ b/mlx/backend/cuda/device/hadamard.cuh @@ -0,0 +1,184 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include "mlx/backend/cuda/device/utils.cuh" + +namespace mlx::core::cu { + +__device__ __forceinline__ void hadamard_radix_m(float* x); + +template +struct Pow2Log2 { + static_assert( + (N > 0) && ((N & (N - 1)) == 0), + "N must be a positive power of two."); + static constexpr int value = 1 + Pow2Log2::value; +}; + +template <> +struct Pow2Log2<1> { + static constexpr int value = 0; +}; + +template +__device__ __forceinline__ void hadamard_radix_pow2(float* x) { + constexpr int kLogR = Pow2Log2::value; + int h = 1; +#pragma unroll + for (int s = 0; s < kLogR; ++s) { +#pragma unroll + for (int i = 0; i < R / 2; ++i) { + int k = i & (h - 1); + int j = ((i - k) << 1) + k; + float a = x[j]; + float b = x[j + h]; + x[j] = a + b; + x[j + h] = a - b; + } + h <<= 1; + } +} + +template +__global__ void +hadamard_n(const T* in, T* out, float scale, long long num_transforms) { + constexpr int kNumThreads = N / max_radix; + constexpr int kLogN = Pow2Log2::value; + constexpr int kLogR = Pow2Log2::value; + constexpr int kNumSteps = kLogN / kLogR; + constexpr int kLogFinal = kLogN % kLogR; + constexpr int kFinalRadix = 1 << kLogFinal; + + if (threadIdx.x >= kNumThreads) { + return; + } + + __shared__ T buf[N]; + int i = threadIdx.x; + + for (long long transform = blockIdx.x; transform < num_transforms; + transform += gridDim.x) { + long long base = (transform / stride) * static_cast(N) * stride + + (transform % stride); + + if constexpr (stride == 1) { +#pragma unroll + for (int j = 0; j < max_radix / read_width; ++j) { + int index = j * read_width * kNumThreads + i * read_width; +#pragma unroll + for (int r = 0; r < read_width; ++r) { + buf[index + r] = in[base + index + r]; + } + } + } else { +#pragma unroll + for (int j = 0; j < max_radix; ++j) { + buf[j * kNumThreads + i] = in[base + (j * kNumThreads + i) * stride]; + } + } + __syncthreads(); + + float x[max_radix]; + int h = 1; + +#pragma unroll + for (int s = 0; s < kNumSteps; ++s) { + int k = i & (h - 1); + int j = ((i - k) << kLogR) + k; + +#pragma unroll + for (int r = 0; r < max_radix; ++r) { + x[r] = static_cast(buf[j + h * r]); + } + + hadamard_radix_pow2(x); + +#pragma unroll + for (int r = 0; r < max_radix; ++r) { + buf[j + h * r] = static_cast(x[r]); + } + + h <<= kLogR; + __syncthreads(); + } + + if constexpr (kFinalRadix > 1) { +#pragma unroll + for (int t = 0; t < max_radix / kFinalRadix; ++t) { + int index = i + t * kNumThreads; + int k = index & (h - 1); + int j = ((index - k) << kLogFinal) + k; +#pragma unroll + for (int r = 0; r < kFinalRadix; ++r) { + x[r] = static_cast(buf[j + h * r]); + } + + hadamard_radix_pow2(x); + +#pragma unroll + for (int r = 0; r < kFinalRadix; ++r) { + buf[j + h * r] = static_cast(x[r]); + } + } + __syncthreads(); + } + + if constexpr (stride == 1) { +#pragma unroll + for (int j = 0; j < max_radix / read_width; ++j) { + int index = j * read_width * kNumThreads + i * read_width; +#pragma unroll + for (int r = 0; r < read_width; ++r) { + float val = static_cast(buf[index + r]); + out[base + index + r] = static_cast(val * scale); + } + } + } else { +#pragma unroll + for (int j = 0; j < max_radix; ++j) { + out[base + (j * kNumThreads + i) * stride] = buf[j * kNumThreads + i]; + } + } + + __syncthreads(); + } +} + +template +__global__ void +hadamard_m(const T* in, T* out, float scale, long long num_tasks) { + constexpr int kTasksPerBatch = N / read_width; + + for (long long task = blockIdx.x * blockDim.x + threadIdx.x; task < num_tasks; + task += blockDim.x * gridDim.x) { + long long i = task % kTasksPerBatch; + long long batch = task / kTasksPerBatch; + long long base = batch * static_cast(M) * N; + + float x[read_width][M]; +#pragma unroll + for (int c = 0; c < M; ++c) { +#pragma unroll + for (int r = 0; r < read_width; ++r) { + x[r][c] = static_cast(in[base + c * N + i * read_width + r]); + } + } + +#pragma unroll + for (int r = 0; r < read_width; ++r) { + hadamard_radix_m(x[r]); + } + +#pragma unroll + for (int c = 0; c < M; ++c) { +#pragma unroll + for (int r = 0; r < read_width; ++r) { + out[base + c * N + i * read_width + r] = + static_cast(x[r][c] * scale); + } + } + } +} + +} // namespace mlx::core::cu diff --git a/mlx/backend/cuda/hadamard.cu b/mlx/backend/cuda/hadamard.cu new file mode 100644 index 00000000..d7def283 --- /dev/null +++ b/mlx/backend/cuda/hadamard.cu @@ -0,0 +1,245 @@ +// Copyright © 2025 Apple Inc. + +#include "mlx/backend/common/hadamard.h" +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/jit_module.h" +#include "mlx/backend/gpu/copy.h" +#include "mlx/dtype_utils.h" +#include "mlx/primitives.h" + +#include +#include + +#include +#include +#include +#include +#include + +namespace mlx::core { + +namespace { + +constexpr int MAX_HADAMARD_THREADS_PER_BLOCK = 256; + +std::string gen_hadamard_codelet(int m) { + std::ostringstream source; + source << "namespace mlx::core::cu {\n"; + source << "__device__ __forceinline__ void hadamard_radix_m(float* x) {\n"; + if (m == 1) { + source << "}\n"; + source << "} // namespace mlx::core::cu\n"; + return source.str(); + } + + auto h_matrices = hadamard_matrices(); + auto it = h_matrices.find(m); + if (it == h_matrices.end()) { + throw std::runtime_error("[hadamard] Invalid radix m."); + } + auto& matrix = it->second; + + source << " float tmp[" << m << "];\n"; + auto start = 1; + auto end = matrix.find('\n', start); + int row_idx = 0; + while (end != std::string_view::npos) { + auto row = matrix.substr(start, end - start); + source << " tmp[" << row_idx << "] ="; + for (int i = 0; i < row.length(); ++i) { + source << " " << row[i] << " x[" << i << "]"; + } + source << ";\n"; + start = end + 1; + end = matrix.find('\n', start); + row_idx++; + } + source << " #pragma unroll\n"; + source << " for (int i = 0; i < " << m << "; ++i) { x[i] = tmp[i]; }\n"; + source << "}\n"; + source << "} // namespace mlx::core::cu\n"; + return source.str(); +} + +std::string hadamard_n_kernel_name( + const Dtype& dtype, + int n, + int max_radix, + int read_width, + int stride) { + return fmt::format( + "mlx::core::cu::hadamard_n<{}, {}, {}, {}, {}>", + dtype_to_cuda_type(dtype), + n, + max_radix, + read_width, + stride); +} + +std::string +hadamard_m_kernel_name(const Dtype& dtype, int n, int m, int read_width) { + return fmt::format( + "mlx::core::cu::hadamard_m<{}, {}, {}, {}>", + dtype_to_cuda_type(dtype), + n, + m, + read_width); +} + +void hadamard_mn_contiguous( + const array& x, + array& y, + int m, + int n1, + int n2, + float scale, + const Stream& s) { + const int n = n1 * n2; + const int read_width_n1 = (n1 == 2) ? 2 : 4; + const int read_width_n2 = (n2 == 2) ? 2 : 4; + const int read_width_m = (n == 2 || m == 28) ? 2 : 4; + const int max_radix_1 = std::min(n1, 16); + const int max_radix_2 = std::min(n2, 16); + const float scale_n1 = 1.0f; + const float scale_n2 = (m == 1) ? scale : 1.0f; + const float scale_m = scale; + + const std::string n1_kernel_name = + hadamard_n_kernel_name(x.dtype(), n1, max_radix_1, read_width_n1, n2); + const std::string n2_kernel_name = + hadamard_n_kernel_name(x.dtype(), n2, max_radix_2, read_width_n2, 1); + const std::string m_kernel_name = + hadamard_m_kernel_name(x.dtype(), n, m, read_width_m); + + const std::string module_name = fmt::format( + "hadamard_{}_{}_{}_{}_{}_{}_{}_{}", + dtype_to_string(x.dtype()), + n, + m, + n1, + n2, + read_width_n1, + read_width_n2, + read_width_m); + + cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() { + std::vector kernel_names = {n2_kernel_name}; + if (n1 > 1) { + kernel_names.push_back(n1_kernel_name); + } + if (m > 1) { + kernel_names.push_back(m_kernel_name); + } + + std::string source = R"( + #include "mlx/backend/cuda/device/utils.cuh" + )"; + source += gen_hadamard_codelet(m); + source += R"( + #include "mlx/backend/cuda/device/hadamard.cuh" + )"; + + return std::make_tuple(false, std::move(source), std::move(kernel_names)); + }); + + auto& encoder = cu::get_command_encoder(s); + + if (n1 > 1) { + const int64_t num_transforms = x.size() / n1; + const uint32_t num_blocks = + static_cast(std::min(num_transforms, 65535)); + + encoder.set_input_array(x); + encoder.set_output_array(y); + + cu::KernelArgs args; + args.append(x); + args.append(y); + args.append(scale_n1); + args.append(num_transforms); + + auto kernel = mod.get_kernel(n1_kernel_name); + encoder.add_kernel_node_raw( + kernel, num_blocks, n1 / max_radix_1, {}, 0, args.args()); + } + + { + const auto& in = (n1 > 1) ? y : x; + const int64_t num_transforms = x.size() / n2; + const uint32_t num_blocks = + static_cast(std::min(num_transforms, 65535)); + + encoder.set_input_array(in); + encoder.set_output_array(y); + + cu::KernelArgs args; + args.append(in); + args.append(y); + args.append(scale_n2); + args.append(num_transforms); + + auto kernel = mod.get_kernel(n2_kernel_name); + encoder.add_kernel_node_raw( + kernel, num_blocks, n2 / max_radix_2, {}, 0, args.args()); + } + + if (m > 1) { + const int64_t num_tasks = x.size() / (m * read_width_m); + const uint32_t block_dim = static_cast( + std::min(num_tasks, MAX_HADAMARD_THREADS_PER_BLOCK)); + const uint32_t num_blocks = static_cast( + std::min((num_tasks + block_dim - 1) / block_dim, 65535)); + + encoder.set_input_array(y); + encoder.set_output_array(y); + + cu::KernelArgs args; + args.append(y); + args.append(y); + args.append(scale_m); + args.append(num_tasks); + + auto kernel = mod.get_kernel(m_kernel_name); + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dim, {}, 0, args.args()); + } +} + +} // namespace + +void Hadamard::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("Hadamard::eval_gpu"); + assert(inputs.size() == 1); + + auto& in = inputs[0]; + if (in.dtype() != float16 && in.dtype() != bfloat16 && + in.dtype() != float32) { + throw std::invalid_argument("[hadamard] Unsupported type."); + } + + // n = m * 2^k where m in (1, 12, 20, 28) + auto [n, m] = decompose_hadamard(in.shape().back()); + int n1 = 1; + int n2 = n; + if (n > 8192) { + for (n2 = 2; n2 * n2 < n; n2 *= 2) { + } + n1 = n / n2; + } + + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + if (in.flags().row_contiguous) { + if (in.is_donatable()) { + out.copy_shared_buffer(in); + } else { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + } + hadamard_mn_contiguous(in, out, m, n1, n2, scale_, s); + } else { + copy_gpu(in, out, CopyType::General, s); + hadamard_mn_contiguous(out, out, m, n1, n2, scale_, s); + } +} + +} // namespace mlx::core diff --git a/mlx/backend/cuda/jit_module.cpp b/mlx/backend/cuda/jit_module.cpp index b0ebb401..953fdb8a 100644 --- a/mlx/backend/cuda/jit_module.cpp +++ b/mlx/backend/cuda/jit_module.cpp @@ -231,6 +231,7 @@ constexpr const char* g_include_names[] = { INCLUDE_PREFIX "config.h", INCLUDE_PREFIX "complex.cuh", INCLUDE_PREFIX "fp16_math.cuh", + INCLUDE_PREFIX "hadamard.cuh", INCLUDE_PREFIX "indexing.cuh", INCLUDE_PREFIX "scatter_ops.cuh", INCLUDE_PREFIX "unary_ops.cuh", @@ -247,6 +248,7 @@ constexpr const char* g_headers[] = { jit_source_config, jit_source_complex, jit_source_fp16_math, + jit_source_hadamard, jit_source_indexing, jit_source_scatter_ops, jit_source_unary_ops, diff --git a/mlx/backend/cuda/primitives.cpp b/mlx/backend/cuda/primitives.cpp index 0caac8de..e29b3e50 100644 --- a/mlx/backend/cuda/primitives.cpp +++ b/mlx/backend/cuda/primitives.cpp @@ -27,7 +27,6 @@ namespace mlx::core { NO_GPU(BlockMaskedMM) NO_GPU(FFT) NO_GPU(GatherQMM) -NO_GPU(Hadamard) NO_GPU_MULTI(LUF) NO_GPU_MULTI(QRF) NO_GPU(SegmentedMM) diff --git a/python/tests/cuda_skip.py b/python/tests/cuda_skip.py index 20793d5c..3b24477b 100644 --- a/python/tests/cuda_skip.py +++ b/python/tests/cuda_skip.py @@ -8,9 +8,6 @@ cuda_skip = { "TestBlas.test_gather_mm_sorted_vjp", # Segmented matmul NYI "TestBlas.test_segmented_mm", - # Hadamard NYI - "TestOps.test_hadamard", - "TestOps.test_hadamard_grad_vmap", # FFTs NYI "TestFFT.test_fft", "TestFFT.test_fft_big_powers_of_two",