chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Remove unused fast-hadamard-transform PyTorch extension sources (#18927)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-02-18 15:51:07 +08:00
committed by GitHub
parent 934b36693c
commit 513c12d23f
2 changed files with 0 additions and 773 deletions
Download Patch File Download Diff File Expand all files Collapse all files

321
python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.cpp
View File

@@ -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 <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <vector>
#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 <typename input_t>
void fast_hadamard_transform_cuda(HadamardParamsBase& params, cudaStream_t stream);
template <typename input_t>
void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, cudaStream_t stream);
template <typename input_t>
void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, cudaStream_t stream);
template <typename input_t>
void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, cudaStream_t stream);
template <typename input_t>
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(&params, 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<input_t>(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<input_t>(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<input_t>(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<input_t>(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<input_t>(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");
}

452
python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_cuda.cu
View File

@@ -1,452 +0,0 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Copied from https://github.com/sgl-project/fast-hadamard-transform
// #pragma once
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
#include <c10/util/BFloat16.h>
#include <c10/util/Half.h>
#include "fast_hadamard_transform.h"
#include "fast_hadamard_transform_common.h"
#include "fast_hadamard_transform_special.h"
#include "static_switch.h"
template <int kNThreads_, int kLogN_, typename input_t_>
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<kNBytes * kNElts>::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 <int kNThreads_, int kLogN_, typename input_t_>
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<kNBytes * kNElts>::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 <int kNThreads_, int kLogN_, typename input_t_>
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<kNBytes * kNElts>::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 <int kNThreads_, int kLogN_, typename input_t_>
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<kNBytes * kNElts>::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 <int kNThreads_, int kLogN_, typename input_t_>
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<kNBytes * kNElts>::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 <int kNChunks>
__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 <int kNChunks>
__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 <int kNChunks>
__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 <int kNChunks>
__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 <typename Ktraits>
__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<vec_t*>(smem_);
const int batch_id = blockIdx.x;
input_t* x = reinterpret_cast<input_t*>(params.x_ptr) + batch_id * params.x_batch_stride;
input_t* out = reinterpret_cast<input_t*>(params.out_ptr) + batch_id * params.out_batch_stride;
float x_vals[kNChunks][kNElts];
load_input<kNChunks, kNElts, input_t>(x, x_vals, params.dim);
hadamard_mult_thread<kLogNElts, kNChunks>(x_vals);
hadamard_mult_warp<kLogWarpSize, 0, kNChunks, kNElts>(x_vals);
if constexpr (kNWarps > 1) {
exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, true, vec_t>(x_vals, smem_exchange);
hadamard_mult_warp<kLogNWarps, 0, kNChunks, kNElts>(x_vals);
exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, false, vec_t>(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<kNElts>(x_vals_transposed);
} else if constexpr (kNChunks == 20) {
hadamard_mult_thread_chunk_20<kNElts>(x_vals_transposed);
} else if constexpr (kNChunks == 28) {
hadamard_mult_thread_chunk_28<kNElts>(x_vals_transposed);
} else if constexpr (kNChunks == 40) {
hadamard_mult_thread_chunk_40<kNElts>(x_vals_transposed);
} else {
constexpr int kLogNChunks = cilog2(kNChunks);
static_assert(1 << kLogNChunks == kNChunks, "kNChunks must be a power of 2");
hadamard_mult_thread<kLogNChunks, kNElts>(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<kNChunks, kNElts, input_t>(out, x_vals, params.dim, params.scale);
}
template <int kNThreads, int kLogN, typename input_t>
void fast_hadamard_transform_launch(HadamardParamsBase& params, cudaStream_t stream) {
using Ktraits = fast_hadamard_transform_kernel_traits<kNThreads, kLogN, input_t>;
constexpr int kSmemSize = Ktraits::kSmemSize;
dim3 grid(params.batch);
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
template <typename input_t>
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 <int kNThreads, int kLogN, typename input_t>
void fast_hadamard_transform_12N_launch(HadamardParamsBase& params, cudaStream_t stream) {
using Ktraits = fast_hadamard_transform_12N_kernel_traits<kNThreads, kLogN, input_t>;
constexpr int kSmemSize = Ktraits::kSmemSize;
dim3 grid(params.batch);
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
template <typename input_t>
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 <int kNThreads, int kLogN, typename input_t>
void fast_hadamard_transform_20N_launch(HadamardParamsBase& params, cudaStream_t stream) {
using Ktraits = fast_hadamard_transform_20N_kernel_traits<kNThreads, kLogN, input_t>;
constexpr int kSmemSize = Ktraits::kSmemSize;
dim3 grid(params.batch);
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
template <typename input_t>
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 <int kNThreads, int kLogN, typename input_t>
void fast_hadamard_transform_28N_launch(HadamardParamsBase& params, cudaStream_t stream) {
using Ktraits = fast_hadamard_transform_28N_kernel_traits<kNThreads, kLogN, input_t>;
constexpr int kSmemSize = Ktraits::kSmemSize;
dim3 grid(params.batch);
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
template <typename input_t>
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 <int kNThreads, int kLogN, typename input_t>
void fast_hadamard_transform_40N_launch(HadamardParamsBase& params, cudaStream_t stream) {
using Ktraits = fast_hadamard_transform_40N_kernel_traits<kNThreads, kLogN, input_t>;
constexpr int kSmemSize = Ktraits::kSmemSize;
dim3 grid(params.batch);
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
template <typename input_t>
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<float>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_12N_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_12N_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_12N_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_20N_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_20N_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_20N_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_28N_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_28N_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_28N_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_40N_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_40N_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
template void fast_hadamard_transform_40N_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);