From 513c12d23f349f62afa02c7b57ebe6f07052e1c5 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Wed, 18 Feb 2026 15:51:07 +0800 Subject: [PATCH] Remove unused fast-hadamard-transform PyTorch extension sources (#18927) --- .../fast_hadamard_transform.cpp | 321 ------------- .../fast_hadamard_transform_cuda.cu | 452 ------------------ 2 files changed, 773 deletions(-) delete mode 100644 python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.cpp delete mode 100644 python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_cuda.cu diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.cpp b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.cpp deleted file mode 100644 index d190f36ca..000000000 --- a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.cpp +++ /dev/null @@ -1,321 +0,0 @@ -/****************************************************************************** - * Copyright (c) 2023, Tri Dao. - ******************************************************************************/ - -// Copied from https://github.com/sgl-project/fast-hadamard-transform - -#include "fast_hadamard_transform.h" - -#include -#include -#include - -#include - -#define CHECK_SHAPE(x, ...) \ - TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")") - -#define DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...) \ - if (ITYPE == at::ScalarType::Half) { \ - using input_t = at::Half; \ - __VA_ARGS__(); \ - } else if (ITYPE == at::ScalarType::BFloat16) { \ - using input_t = at::BFloat16; \ - __VA_ARGS__(); \ - } else if (ITYPE == at::ScalarType::Float) { \ - using input_t = float; \ - __VA_ARGS__(); \ - } else { \ - AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \ - } - -template -void fast_hadamard_transform_cuda(HadamardParamsBase& params, cudaStream_t stream); - -template -void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, cudaStream_t stream); - -template -void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, cudaStream_t stream); - -template -void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, cudaStream_t stream); - -template -void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, cudaStream_t stream); - -void set_hadamard_params( - HadamardParamsBase& params, - // sizes - const size_t batch, - const size_t dim, - const size_t multiple, - // device pointers - const at::Tensor x, - const at::Tensor out, - float scale) { - // Reset the parameters - memset(¶ms, 0, sizeof(params)); - - params.batch = batch; - params.dim = dim; - params.log_N = int(ceil(std::log2(dim / multiple))); - - // Set the pointers and strides. - params.x_ptr = x.data_ptr(); - params.out_ptr = out.data_ptr(); - // All stride are in elements, not bytes. - params.x_batch_stride = x.stride(0); - params.out_batch_stride = out.stride(0); - - params.scale = scale; -} - -at::Tensor fast_hadamard_transform(at::Tensor& x, float scale) { - auto input_type = x.scalar_type(); - TORCH_CHECK( - input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || - input_type == at::ScalarType::BFloat16); - - TORCH_CHECK(x.is_cuda()); - - const auto shapes_og = x.sizes(); - const int dim_og = x.size(-1); - x = x.reshape({-1, dim_og}); - if (x.stride(-1) != 1) { - x = x.contiguous(); - } - const auto sizes = x.sizes(); - const int batch_size = sizes[0]; - - CHECK_SHAPE(x, batch_size, dim_og); - TORCH_CHECK(x.stride(1) == 1); - - if (dim_og % 8 != 0) { - x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 8 - dim_og % 8})); - } - const int dim = x.size(1); - - TORCH_CHECK(dim % 8 == 0, "fast_hadamard_transform only supports hidden dimension divisible by 8 for now"); - TORCH_CHECK(dim <= 32768, "fast_hadamard_transform only supports hidden dimension at most 32768 for now"); - - at::Tensor out = torch::empty_like(x); - - HadamardParamsBase params; - set_hadamard_params(params, batch_size, dim, 1, x, out, scale); - - // Otherwise the kernel will be launched from cuda:0 device - // Cast to char to avoid compiler warning about narrowing - at::cuda::CUDAGuard device_guard{(char)x.get_device()}; - auto stream = at::cuda::getCurrentCUDAStream().stream(); - DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16( - x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_cuda(params, stream); }); - if (dim_og % 8 != 0) { - out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)}); - } - return out.reshape(shapes_og); -} - -at::Tensor fast_hadamard_transform_12N(at::Tensor& x, float scale) { - auto input_type = x.scalar_type(); - TORCH_CHECK( - input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || - input_type == at::ScalarType::BFloat16); - - TORCH_CHECK(x.is_cuda()); - - const auto shapes_og = x.sizes(); - const int dim_og = x.size(-1); - x = x.reshape({-1, dim_og}); - if (x.stride(-1) != 1) { - x = x.contiguous(); - } - const auto sizes = x.sizes(); - const int batch_size = sizes[0]; - - CHECK_SHAPE(x, batch_size, dim_og); - TORCH_CHECK(x.stride(1) == 1); - - if (dim_og % (4 * 12) != 0) { - x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, (4 * 12) - dim_og % (4 * 12)})); - } - const int dim = x.size(1); - - TORCH_CHECK( - dim % (4 * 12) == 0, "fast_hadamard_transform_12N only supports hidden dimension divisible by 48 for now"); - TORCH_CHECK(dim <= 12 * 1024, "fast_hadamard_transform_12N only supports hidden dimension at most 12288 for now"); - - at::Tensor out = torch::empty_like(x); - - HadamardParamsBase params; - set_hadamard_params(params, batch_size, dim, 12, x, out, scale); - - // Otherwise the kernel will be launched from cuda:0 device - // Cast to char to avoid compiler warning about narrowing - at::cuda::CUDAGuard device_guard{(char)x.get_device()}; - auto stream = at::cuda::getCurrentCUDAStream().stream(); - DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16( - x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_12N_cuda(params, stream); }); - if (dim_og % (4 * 12) != 0) { - out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)}); - } - return out.reshape(shapes_og); -} - -at::Tensor fast_hadamard_transform_20N(at::Tensor& x, float scale) { - auto input_type = x.scalar_type(); - TORCH_CHECK( - input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || - input_type == at::ScalarType::BFloat16); - - TORCH_CHECK(x.is_cuda()); - - const auto shapes_og = x.sizes(); - const int dim_og = x.size(-1); - x = x.reshape({-1, dim_og}); - if (x.stride(-1) != 1) { - x = x.contiguous(); - } - const auto sizes = x.sizes(); - const int batch_size = sizes[0]; - - CHECK_SHAPE(x, batch_size, dim_og); - TORCH_CHECK(x.stride(1) == 1); - - if (dim_og % (4 * 20) != 0) { - x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, (4 * 20) - dim_og % (4 * 20)})); - } - const int dim = x.size(1); - - TORCH_CHECK( - dim % (4 * 20) == 0, "fast_hadamard_transform_20N only supports hidden dimension divisible by 80 for now"); - TORCH_CHECK(dim <= 20 * 1024, "fast_hadamard_transform_20N only supports hidden dimension at most 20480 for now"); - - at::Tensor out = torch::empty_like(x); - - HadamardParamsBase params; - set_hadamard_params(params, batch_size, dim, 20, x, out, scale); - - // Otherwise the kernel will be launched from cuda:0 device - // Cast to char to avoid compiler warning about narrowing - at::cuda::CUDAGuard device_guard{(char)x.get_device()}; - auto stream = at::cuda::getCurrentCUDAStream().stream(); - DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16( - x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_20N_cuda(params, stream); }); - if (dim_og % (4 * 20) != 0) { - out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)}); - } - return out.reshape(shapes_og); -} - -at::Tensor fast_hadamard_transform_28N(at::Tensor& x, float scale) { - auto input_type = x.scalar_type(); - TORCH_CHECK( - input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || - input_type == at::ScalarType::BFloat16); - - TORCH_CHECK(x.is_cuda()); - - const auto shapes_og = x.sizes(); - const int dim_og = x.size(-1); - x = x.reshape({-1, dim_og}); - if (x.stride(-1) != 1) { - x = x.contiguous(); - } - const auto sizes = x.sizes(); - const int batch_size = sizes[0]; - - CHECK_SHAPE(x, batch_size, dim_og); - TORCH_CHECK(x.stride(1) == 1); - - if (dim_og % (4 * 28) != 0) { - x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, (4 * 28) - dim_og % (4 * 28)})); - } - const int dim = x.size(1); - - TORCH_CHECK( - dim % (4 * 28) == 0, "fast_hadamard_transform_28N only supports hidden dimension divisible by 112 for now"); - TORCH_CHECK(dim <= 28 * 1024, "fast_hadamard_transform_28N only supports hidden dimension at most 28672 for now"); - - at::Tensor out = torch::empty_like(x); - - HadamardParamsBase params; - set_hadamard_params(params, batch_size, dim, 28, x, out, scale); - - // Otherwise the kernel will be launched from cuda:0 device - // Cast to char to avoid compiler warning about narrowing - at::cuda::CUDAGuard device_guard{(char)x.get_device()}; - auto stream = at::cuda::getCurrentCUDAStream().stream(); - DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16( - x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_28N_cuda(params, stream); }); - if (dim_og % (8 * 28) != 0) { - out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)}); - } - return out.reshape(shapes_og); -} - -at::Tensor fast_hadamard_transform_40N(at::Tensor& x, float scale) { - auto input_type = x.scalar_type(); - TORCH_CHECK( - input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || - input_type == at::ScalarType::BFloat16); - - TORCH_CHECK(x.is_cuda()); - - const auto shapes_og = x.sizes(); - const int dim_og = x.size(-1); - x = x.reshape({-1, dim_og}); - if (x.stride(-1) != 1) { - x = x.contiguous(); - } - const auto sizes = x.sizes(); - const int batch_size = sizes[0]; - - CHECK_SHAPE(x, batch_size, dim_og); - TORCH_CHECK(x.stride(1) == 1); - - if (dim_og % (4 * 40) != 0) { - x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, (4 * 40) - dim_og % (4 * 40)})); - } - const int dim = x.size(1); - - TORCH_CHECK( - dim % (4 * 40) == 0, "fast_hadamard_transform_40N only supports hidden dimension divisible by 160 for now"); - TORCH_CHECK(dim <= 40 * 1024, "fast_hadamard_transform_40N only supports hidden dimension at most 40960 for now"); - - at::Tensor out = torch::empty_like(x); - - HadamardParamsBase params; - set_hadamard_params(params, batch_size, dim, 40, x, out, scale); - - // Otherwise the kernel will be launched from cuda:0 device - // Cast to char to avoid compiler warning about narrowing - at::cuda::CUDAGuard device_guard{(char)x.get_device()}; - auto stream = at::cuda::getCurrentCUDAStream().stream(); - DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16( - x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_40N_cuda(params, stream); }); - if (dim_og % (8 * 40) != 0) { - out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)}); - } - return out.reshape(shapes_og); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("fast_hadamard_transform", &fast_hadamard_transform, "Fast Hadamard transform"); - m.def( - "fast_hadamard_transform_12N", - &fast_hadamard_transform_12N, - "Fast Hadamard transform with dimension = 12 * power of 2"); - m.def( - "fast_hadamard_transform_20N", - &fast_hadamard_transform_20N, - "Fast Hadamard transform with dimension = 20 * power of 2"); - m.def( - "fast_hadamard_transform_28N", - &fast_hadamard_transform_28N, - "Fast Hadamard transform with dimension = 28 * power of 2"); - m.def( - "fast_hadamard_transform_40N", - &fast_hadamard_transform_40N, - "Fast Hadamard transform with dimension = 40 * power of 2"); -} diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_cuda.cu b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_cuda.cu deleted file mode 100644 index 6c3fcb0ff..000000000 --- a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_cuda.cu +++ /dev/null @@ -1,452 +0,0 @@ -/****************************************************************************** - * Copyright (c) 2023, Tri Dao. - ******************************************************************************/ - -// Copied from https://github.com/sgl-project/fast-hadamard-transform - -// #pragma once - -#include // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK -#include -#include - -#include "fast_hadamard_transform.h" -#include "fast_hadamard_transform_common.h" -#include "fast_hadamard_transform_special.h" -#include "static_switch.h" - -template -struct fast_hadamard_transform_kernel_traits { - using input_t = input_t_; - static constexpr int kNThreads = kNThreads_; - static constexpr int kLogN = kLogN_; - static constexpr int N = 1 << kLogN; - static constexpr int kNBytes = sizeof(input_t); - static_assert(kNBytes == 2 || kNBytes == 4); - static constexpr int kNElts = kNBytes == 4 ? 4 : 8; - // It's possible that we need to do 2 rounds of exchange if input_t is 16 bits - // (since then we'd have 8 values of float, and each round we can exchange 4 floats). - static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t); - using vec_t = typename BytesToType::Type; - static constexpr int kNChunks = N / (kNElts * kNThreads); - // We don't want to use more than 32 KB of shared memory. - static constexpr int kSmemExchangeSize = std::min(N * 4, 32 * 1024); - static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; - static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); - static constexpr int kSmemSize = kSmemExchangeSize; -}; - -template -struct fast_hadamard_transform_12N_kernel_traits { - using input_t = input_t_; - static constexpr int kNThreads = kNThreads_; - static constexpr int kLogN = kLogN_; - static constexpr int N = (1 << kLogN) * 12; - static_assert(N <= 12 * 1024, "fast_hadamard_transform_12 only supports dim <= 12288"); - static constexpr int kNBytes = sizeof(input_t); - static_assert(kNBytes == 2 || kNBytes == 4); - static constexpr int kNElts = 4; - // It's possible that we need to do 2 rounds of exchange if input_t is 16 bits - // (since then we'd have 8 values of float, and each round we can exchange 4 floats). - static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t); - using vec_t = typename BytesToType::Type; - static constexpr int kNChunks = N / (kNElts * kNThreads); - static_assert(kNChunks == 12); - // We don't want to use more than 24 KB of shared memory. - static constexpr int kSmemExchangeSize = std::min(N * 4, 24 * 1024); - static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; - static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); - static constexpr int kSmemSize = kSmemExchangeSize; -}; - -template -struct fast_hadamard_transform_20N_kernel_traits { - using input_t = input_t_; - static constexpr int kNThreads = kNThreads_; - static constexpr int kLogN = kLogN_; - static constexpr int N = (1 << kLogN) * 20; - static_assert(N <= 20 * 1024, "fast_hadamard_transform_20 only supports dim <= 20480"); - static constexpr int kNBytes = sizeof(input_t); - static_assert(kNBytes == 2 || kNBytes == 4); - static constexpr int kNElts = 4; - // It's possible that we need to do 2 rounds of exchange if input_t is 16 bits - // (since then we'd have 8 values of float, and each round we can exchange 4 floats). - static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t); - using vec_t = typename BytesToType::Type; - static constexpr int kNChunks = N / (kNElts * kNThreads); - static_assert(kNChunks == 20); - // We don't want to use more than 40 KB of shared memory. - static constexpr int kSmemExchangeSize = std::min(N * 4, 40 * 1024); - static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; - static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); - static constexpr int kSmemSize = kSmemExchangeSize; -}; - -template -struct fast_hadamard_transform_28N_kernel_traits { - using input_t = input_t_; - static constexpr int kNThreads = kNThreads_; - static constexpr int kLogN = kLogN_; - static constexpr int N = (1 << kLogN) * 28; - static_assert(N <= 28 * 1024, "fast_hadamard_transform_28 only supports dim <= 28672"); - static constexpr int kNBytes = sizeof(input_t); - static_assert(kNBytes == 2 || kNBytes == 4); - static constexpr int kNElts = 4; - // It's possible that we need to do 2 rounds of exchange if input_t is 16 bits - // (since then we'd have 8 values of float, and each round we can exchange 4 floats). - static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t); - using vec_t = typename BytesToType::Type; - static constexpr int kNChunks = N / (kNElts * kNThreads); - static_assert(kNChunks == 28); - // We don't want to use more than 28 KB of shared memory. - static constexpr int kSmemExchangeSize = std::min(N * 4, 28 * 1024); - static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; - static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); - static constexpr int kSmemSize = kSmemExchangeSize; -}; - -template -struct fast_hadamard_transform_40N_kernel_traits { - using input_t = input_t_; - static constexpr int kNThreads = kNThreads_; - static constexpr int kLogN = kLogN_; - static constexpr int N = (1 << kLogN) * 40; - static_assert(N <= 40 * 1024, "fast_hadamard_transform_40 only supports dim <= 40960"); - static constexpr int kNBytes = sizeof(input_t); - static_assert(kNBytes == 2 || kNBytes == 4); - static constexpr int kNElts = 4; - // It's possible that we need to do 2 rounds of exchange if input_t is 16 bits - // (since then we'd have 8 values of float, and each round we can exchange 4 floats). - static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t); - using vec_t = typename BytesToType::Type; - static constexpr int kNChunks = N / (kNElts * kNThreads); - static_assert(kNChunks == 40); - // We don't want to use more than 40 KB of shared memory. - static constexpr int kSmemExchangeSize = std::min(N * 4, 40 * 1024); - static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; - static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); - static constexpr int kSmemSize = kSmemExchangeSize; -}; - -template -__device__ __forceinline__ void hadamard_mult_thread_chunk_12(float x[kNChunks][12]) { -#pragma unroll - for (int c = 0; c < kNChunks; ++c) { - hadamard_mult_thread_12(x[c]); - } -} - -template -__device__ __forceinline__ void hadamard_mult_thread_chunk_20(float x[kNChunks][20]) { -#pragma unroll - for (int c = 0; c < kNChunks; ++c) { - hadamard_mult_thread_20(x[c]); - } -} - -template -__device__ __forceinline__ void hadamard_mult_thread_chunk_28(float x[kNChunks][28]) { -#pragma unroll - for (int c = 0; c < kNChunks; ++c) { - hadamard_mult_thread_28(x[c]); - } -} - -template -__device__ __forceinline__ void hadamard_mult_thread_chunk_40(float x[kNChunks][40]) { -#pragma unroll - for (int c = 0; c < kNChunks; ++c) { - hadamard_mult_thread_40(x[c]); - } -} - -template -__global__ __launch_bounds__(Ktraits::kNThreads) void fast_hadamard_transform_kernel(HadamardParamsBase params) { - constexpr int kNThreads = Ktraits::kNThreads; - constexpr int kNElts = Ktraits::kNElts; - constexpr int kNExchangePerVec = Ktraits::kNExchangePerVec; - constexpr int kNExchangeRounds = Ktraits::kNExchangeRounds; - constexpr int kNChunks = Ktraits::kNChunks; - using input_t = typename Ktraits::input_t; - using vec_t = typename Ktraits::vec_t; - - constexpr int kLogNElts = cilog2(Ktraits::kNElts); - static_assert(1 << kLogNElts == kNElts, "kNElts must be a power of 2"); - constexpr int kWarpSize = std::min(kNThreads, 32); - constexpr int kLogWarpSize = cilog2(kWarpSize); - static_assert(1 << kLogWarpSize == kWarpSize, "Warp size must be a power of 2"); - constexpr int kNWarps = kNThreads / kWarpSize; - constexpr int kLogNWarps = cilog2(kNWarps); - static_assert(1 << kLogNWarps == kNWarps, "kNWarps must be a power of 2"); - constexpr int kLoadsPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNThreads); - static_assert( - kLoadsPerExchange * sizeof(vec_t) * kNThreads == Ktraits::kSmemExchangeSize, - "kSmemExchangeSize should be a power of 2"); - static_assert(kNExchangeRounds * kLoadsPerExchange * sizeof(vec_t) == kNChunks * kNElts * sizeof(float)); - - constexpr int kChunksPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNExchangePerVec * kNThreads); - static_assert(kChunksPerExchange * sizeof(vec_t) * kNExchangePerVec * kNThreads == Ktraits::kSmemExchangeSize); - constexpr int kNExchanges = kNChunks / kChunksPerExchange; - static_assert(kNExchanges * kChunksPerExchange == kNChunks); - - // Shared memory. - extern __shared__ char smem_[]; - vec_t* smem_exchange = reinterpret_cast(smem_); - - const int batch_id = blockIdx.x; - input_t* x = reinterpret_cast(params.x_ptr) + batch_id * params.x_batch_stride; - input_t* out = reinterpret_cast(params.out_ptr) + batch_id * params.out_batch_stride; - - float x_vals[kNChunks][kNElts]; - load_input(x, x_vals, params.dim); - - hadamard_mult_thread(x_vals); - hadamard_mult_warp(x_vals); - - if constexpr (kNWarps > 1) { - exchange_smem_pre(x_vals, smem_exchange); - hadamard_mult_warp(x_vals); - exchange_smem_pre(x_vals, smem_exchange); - } - - if constexpr (kNChunks > 1) { - float x_vals_transposed[kNElts][kNChunks]; -#pragma unroll - for (int c = 0; c < kNChunks; ++c) { -#pragma unroll - for (int i = 0; i < kNElts; ++i) { - x_vals_transposed[i][c] = x_vals[c][i]; - } - } - if constexpr (kNChunks == 12) { - hadamard_mult_thread_chunk_12(x_vals_transposed); - } else if constexpr (kNChunks == 20) { - hadamard_mult_thread_chunk_20(x_vals_transposed); - } else if constexpr (kNChunks == 28) { - hadamard_mult_thread_chunk_28(x_vals_transposed); - } else if constexpr (kNChunks == 40) { - hadamard_mult_thread_chunk_40(x_vals_transposed); - } else { - constexpr int kLogNChunks = cilog2(kNChunks); - static_assert(1 << kLogNChunks == kNChunks, "kNChunks must be a power of 2"); - hadamard_mult_thread(x_vals_transposed); - } -#pragma unroll - for (int c = 0; c < kNChunks; ++c) { -#pragma unroll - for (int i = 0; i < kNElts; ++i) { - x_vals[c][i] = x_vals_transposed[i][c]; - } - } - } - - store_output(out, x_vals, params.dim, params.scale); -} - -template -void fast_hadamard_transform_launch(HadamardParamsBase& params, cudaStream_t stream) { - using Ktraits = fast_hadamard_transform_kernel_traits; - constexpr int kSmemSize = Ktraits::kSmemSize; - dim3 grid(params.batch); - auto kernel = &fast_hadamard_transform_kernel; - if (kSmemSize >= 48 * 1024) { - C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); - } - kernel<<>>(params); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -template -void fast_hadamard_transform_cuda(HadamardParamsBase& params, cudaStream_t stream) { - if (params.log_N == 3) { - fast_hadamard_transform_launch<1, 3, input_t>(params, stream); - } else if (params.log_N == 4) { - fast_hadamard_transform_launch<2, 4, input_t>(params, stream); - } else if (params.log_N == 5) { - fast_hadamard_transform_launch<4, 5, input_t>(params, stream); - } else if (params.log_N == 6) { - fast_hadamard_transform_launch<8, 6, input_t>(params, stream); - } else if (params.log_N == 7) { - fast_hadamard_transform_launch<16, 7, input_t>(params, stream); - } else if (params.log_N == 8) { - fast_hadamard_transform_launch<32, 8, input_t>(params, stream); - } else if (params.log_N == 9) { - fast_hadamard_transform_launch<32, 9, input_t>(params, stream); - } else if (params.log_N == 10) { - fast_hadamard_transform_launch<128, 10, input_t>(params, stream); - } else if (params.log_N == 11) { - fast_hadamard_transform_launch<256, 11, input_t>(params, stream); - } else if (params.log_N == 12) { - fast_hadamard_transform_launch<256, 12, input_t>(params, stream); - } else if (params.log_N == 13) { - fast_hadamard_transform_launch<256, 13, input_t>(params, stream); - } else if (params.log_N == 14) { - fast_hadamard_transform_launch<256, 14, input_t>(params, stream); - } else if (params.log_N == 15) { - fast_hadamard_transform_launch<256, 15, input_t>(params, stream); - } -} - -template -void fast_hadamard_transform_12N_launch(HadamardParamsBase& params, cudaStream_t stream) { - using Ktraits = fast_hadamard_transform_12N_kernel_traits; - constexpr int kSmemSize = Ktraits::kSmemSize; - dim3 grid(params.batch); - auto kernel = &fast_hadamard_transform_kernel; - if (kSmemSize >= 48 * 1024) { - C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); - } - kernel<<>>(params); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -template -void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, cudaStream_t stream) { - if (params.log_N == 2) { - fast_hadamard_transform_12N_launch<1, 2, input_t>(params, stream); - } else if (params.log_N == 3) { - fast_hadamard_transform_12N_launch<2, 3, input_t>(params, stream); - } else if (params.log_N == 4) { - fast_hadamard_transform_12N_launch<4, 4, input_t>(params, stream); - } else if (params.log_N == 5) { - fast_hadamard_transform_12N_launch<8, 5, input_t>(params, stream); - } else if (params.log_N == 6) { - fast_hadamard_transform_12N_launch<16, 6, input_t>(params, stream); - } else if (params.log_N == 7) { - fast_hadamard_transform_12N_launch<32, 7, input_t>(params, stream); - } else if (params.log_N == 8) { - fast_hadamard_transform_12N_launch<64, 8, input_t>(params, stream); - } else if (params.log_N == 9) { - fast_hadamard_transform_12N_launch<128, 9, input_t>(params, stream); - } else if (params.log_N == 10) { - fast_hadamard_transform_12N_launch<256, 10, input_t>(params, stream); - } -} - -template -void fast_hadamard_transform_20N_launch(HadamardParamsBase& params, cudaStream_t stream) { - using Ktraits = fast_hadamard_transform_20N_kernel_traits; - constexpr int kSmemSize = Ktraits::kSmemSize; - dim3 grid(params.batch); - auto kernel = &fast_hadamard_transform_kernel; - if (kSmemSize >= 48 * 1024) { - C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); - } - kernel<<>>(params); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -template -void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, cudaStream_t stream) { - if (params.log_N == 2) { - fast_hadamard_transform_20N_launch<1, 2, input_t>(params, stream); - } else if (params.log_N == 3) { - fast_hadamard_transform_20N_launch<2, 3, input_t>(params, stream); - } else if (params.log_N == 4) { - fast_hadamard_transform_20N_launch<4, 4, input_t>(params, stream); - } else if (params.log_N == 5) { - fast_hadamard_transform_20N_launch<8, 5, input_t>(params, stream); - } else if (params.log_N == 6) { - fast_hadamard_transform_20N_launch<16, 6, input_t>(params, stream); - } else if (params.log_N == 7) { - fast_hadamard_transform_20N_launch<32, 7, input_t>(params, stream); - } else if (params.log_N == 8) { - fast_hadamard_transform_20N_launch<64, 8, input_t>(params, stream); - } else if (params.log_N == 9) { - fast_hadamard_transform_20N_launch<128, 9, input_t>(params, stream); - } else if (params.log_N == 10) { - fast_hadamard_transform_20N_launch<256, 10, input_t>(params, stream); - } -} - -template -void fast_hadamard_transform_28N_launch(HadamardParamsBase& params, cudaStream_t stream) { - using Ktraits = fast_hadamard_transform_28N_kernel_traits; - constexpr int kSmemSize = Ktraits::kSmemSize; - dim3 grid(params.batch); - auto kernel = &fast_hadamard_transform_kernel; - if (kSmemSize >= 48 * 1024) { - C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); - } - kernel<<>>(params); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -template -void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, cudaStream_t stream) { - if (params.log_N == 2) { - fast_hadamard_transform_28N_launch<1, 2, input_t>(params, stream); - } else if (params.log_N == 3) { - fast_hadamard_transform_28N_launch<2, 3, input_t>(params, stream); - } else if (params.log_N == 4) { - fast_hadamard_transform_28N_launch<4, 4, input_t>(params, stream); - } else if (params.log_N == 5) { - fast_hadamard_transform_28N_launch<8, 5, input_t>(params, stream); - } else if (params.log_N == 6) { - fast_hadamard_transform_28N_launch<16, 6, input_t>(params, stream); - } else if (params.log_N == 7) { - fast_hadamard_transform_28N_launch<32, 7, input_t>(params, stream); - } else if (params.log_N == 8) { - fast_hadamard_transform_28N_launch<64, 8, input_t>(params, stream); - } else if (params.log_N == 9) { - fast_hadamard_transform_28N_launch<128, 9, input_t>(params, stream); - } else if (params.log_N == 10) { - fast_hadamard_transform_28N_launch<256, 10, input_t>(params, stream); - } -} - -template -void fast_hadamard_transform_40N_launch(HadamardParamsBase& params, cudaStream_t stream) { - using Ktraits = fast_hadamard_transform_40N_kernel_traits; - constexpr int kSmemSize = Ktraits::kSmemSize; - dim3 grid(params.batch); - auto kernel = &fast_hadamard_transform_kernel; - if (kSmemSize >= 48 * 1024) { - C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); - } - kernel<<>>(params); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -template -void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, cudaStream_t stream) { - if (params.log_N == 2) { - fast_hadamard_transform_40N_launch<1, 2, input_t>(params, stream); - } else if (params.log_N == 3) { - fast_hadamard_transform_40N_launch<2, 3, input_t>(params, stream); - } else if (params.log_N == 4) { - fast_hadamard_transform_40N_launch<4, 4, input_t>(params, stream); - } else if (params.log_N == 5) { - fast_hadamard_transform_40N_launch<8, 5, input_t>(params, stream); - } else if (params.log_N == 6) { - fast_hadamard_transform_40N_launch<16, 6, input_t>(params, stream); - } else if (params.log_N == 7) { - fast_hadamard_transform_40N_launch<32, 7, input_t>(params, stream); - } else if (params.log_N == 8) { - fast_hadamard_transform_40N_launch<64, 8, input_t>(params, stream); - } else if (params.log_N == 9) { - fast_hadamard_transform_40N_launch<128, 9, input_t>(params, stream); - } else if (params.log_N == 10) { - fast_hadamard_transform_40N_launch<256, 10, input_t>(params, stream); - } -} - -template void fast_hadamard_transform_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_cuda(HadamardParamsBase& params, cudaStream_t stream); - -template void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, cudaStream_t stream); - -template void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, cudaStream_t stream); - -template void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, cudaStream_t stream); - -template void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, cudaStream_t stream); -template void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, cudaStream_t stream);