diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py new file mode 100644 index 000000000..b19a8ba69 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/code_gen.py @@ -0,0 +1,197 @@ +from pathlib import Path + +import numpy as np + +# From https://en.wikipedia.org/wiki/Paley_construction (construction II for q = 5) + +had_12_paley = """ ++-++++++++++ +--+-+-+-+-+- ++++-++----++ ++---+--+-++- ++++++-++---- ++-+---+--+-+ +++--+++-++-- ++--++---+--+ +++----+++-++ ++--+-++---+- +++++----+++- ++-+--+-++--- +""" + +# From http://neilsloane.com/hadamard/ + +had_12 = """ ++----------- +++-+---+++-+ ++++-+---+++- ++-++-+---+++ +++-++-+---++ ++++-++-+---+ +++++-++-+--- ++-+++-++-+-- ++--+++-++-+- ++---+++-++-+ +++---+++-++- ++-+---+++-++ +""" + +had_20_will = """ ++----+----++--++-++- +-+----+---+++---+-++ +--+----+---+++-+-+-+ +---+----+---+++++-+- +----+----++--++-++-+ +-+++++-----+--+++--+ ++-+++-+---+-+--+++-- +++-++--+---+-+--+++- ++++-+---+---+-+--+++ +++++-----++--+-+--++ +--++-+-++-+-----++++ +---++-+-++-+---+-+++ ++---++-+-+--+--++-++ +++---++-+----+-+++-+ +-++---++-+----+++++- +-+--+--++-+----+---- ++-+-----++-+----+--- +-+-+-+---+--+----+-- +--+-+++------+----+- ++--+--++------+----+ +""" + + +had_28_will = """ ++------++----++-+--+-+--++-- +-+-----+++-----+-+--+-+--++- +--+-----+++---+-+-+----+--++ +---+-----+++---+-+-+-+--+--+ +----+-----+++---+-+-+++--+-- +-----+-----++++--+-+--++--+- +------++----++-+--+-+--++--+ +--++++-+-------++--+++-+--+- +---++++-+-----+-++--+-+-+--+ ++---+++--+----++-++--+-+-+-- +++---++---+----++-++--+-+-+- ++++---+----+----++-++--+-+-+ +++++--------+-+--++-++--+-+- +-++++--------+++--++--+--+-+ +-+-++-++--++--+--------++++- ++-+-++--+--++--+--------++++ +-+-+-++--+--++--+----+---+++ ++-+-+-++--+--+---+---++---++ +++-+-+-++--+------+--+++---+ +-++-+-+-++--+------+-++++--- ++-++-+---++--+------+-++++-- +-++--++-+-++-+++----++------ ++-++--++-+-++-+++-----+----- +++-++---+-+-++-+++-----+---- +-++-++-+-+-+-+--+++-----+--- +--++-++++-+-+----+++-----+-- ++--++-+-++-+-+----+++-----+- +++--++-+-++-+-+----++------+ +""" + + +had_40_tpal = """ ++-------------------+------------------- +++-++----+-+-++++--+++-++----+-+-++++--+ ++++-++----+-+-++++--+++-++----+-+-++++-- ++-++-++----+-+-++++-+-++-++----+-+-++++- ++--++-++----+-+-+++++--++-++----+-+-++++ +++--++-++----+-+-+++++--++-++----+-+-+++ ++++--++-++----+-+-+++++--++-++----+-+-++ +++++--++-++----+-+-+++++--++-++----+-+-+ ++++++--++-++----+-+-+++++--++-++----+-+- ++-++++--++-++----+-++-++++--++-++----+-+ +++-++++--++-++----+-++-++++--++-++----+- ++-+-++++--++-++----++-+-++++--++-++----+ +++-+-++++--++-++----++-+-++++--++-++---- ++-+-+-++++--++-++---+-+-+-++++--++-++--- ++--+-+-++++--++-++--+--+-+-++++--++-++-- ++---+-+-++++--++-++-+---+-+-++++--++-++- ++----+-+-++++--++-+++----+-+-++++--++-++ +++----+-+-++++--++-+++----+-+-++++--++-+ ++++----+-+-++++--++-+++----+-+-++++--++- ++-++----+-+-++++--+++-++----+-+-++++--++ ++--------------------+++++++++++++++++++ +++-++----+-+-++++--+--+--++++-+-+----++- ++++-++----+-+-++++-----+--++++-+-+----++ ++-++-++----+-+-++++--+--+--++++-+-+----+ ++--++-++----+-+-++++-++--+--++++-+-+---- +++--++-++----+-+-+++--++--+--++++-+-+--- ++++--++-++----+-+-++---++--+--++++-+-+-- +++++--++-++----+-+-+----++--+--++++-+-+- ++++++--++-++----+-+------++--+--++++-+-+ ++-++++--++-++----+-+-+----++--+--++++-+- +++-++++--++-++----+---+----++--+--++++-+ ++-+-++++--++-++----+-+-+----++--+--++++- +++-+-++++--++-++------+-+----++--+--++++ ++-+-+-++++--++-++----+-+-+----++--+--+++ ++--+-+-++++--++-++---++-+-+----++--+--++ ++---+-+-++++--++-++--+++-+-+----++--+--+ ++----+-+-++++--++-++-++++-+-+----++--+-- +++----+-+-++++--++-+--++++-+-+----++--+- ++++----+-+-++++--++----++++-+-+----++--+ ++-++----+-+-++++--++-+--++++-+-+----++-- +""" + + +header = """ +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ + +// This file is auto-generated. See "code_gen.py"\n + +#pragma once + +""" + +template = """ +__device__ __forceinline__ void hadamard_mult_thread_{N}(float x[{N}]) { + float out[{N}]; + {code} + #pragma unroll + for (int i = 0; i < {N}; i++) { x[i] = out[i]; } +} + +""" + + +def string_to_array(string): + # Convert strings of + and - to bool arrays + string = string.strip().replace("+", "1").replace("-", "-1").split() + return np.stack( + [ + np.fromstring(" ".join(string[i]), dtype=np.int32, sep=" ") + for i in range(len(string)) + ] + ) + + +def array_code_gen(arr): + N = arr.shape[0] + assert arr.shape[0] == arr.shape[1] + out = [] + for i in range(N): + out.append( + f"out[{i}] = " + + " ".join([f"{'+' if arr[i, j] == 1 else '-'} x[{j}]" for j in range(N)]) + + ";" + ) + return template.replace("{N}", str(N)).replace("{code}", "\n ".join(out)) + + +def main(): + output_dir = Path(__file__).parent / "fast_hadamard_transform_special.h" + output_dir.write_text( + header + + array_code_gen(string_to_array(had_12_paley)) + + array_code_gen(string_to_array(had_20_will)) + + array_code_gen(string_to_array(had_28_will)) + + array_code_gen(string_to_array(had_40_tpal)) + ) + + +if __name__ == "__main__": + main() 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 new file mode 100644 index 000000000..d190f36ca --- /dev/null +++ b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.cpp @@ -0,0 +1,321 @@ +/****************************************************************************** + * 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.h b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.h new file mode 100644 index 000000000..1dda51c3e --- /dev/null +++ b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform.h @@ -0,0 +1,24 @@ +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ + +// Copied from https://github.com/sgl-project/fast-hadamard-transform + +#pragma once + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct HadamardParamsBase { + using index_t = int64_t; + + int batch, dim, log_N; + + index_t x_batch_stride; + index_t out_batch_stride; + + float scale; + + // Common data pointers. + void* __restrict__ x_ptr; + void* __restrict__ out_ptr; +}; diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h new file mode 100644 index 000000000..f6e6117d5 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_common.h @@ -0,0 +1,214 @@ +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ + +// Copied from https://github.com/sgl-project/fast-hadamard-transform + +#pragma once + +#include +#include + +#define FULL_MASK 0xffffffff + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct uint8 { + uint4 u; + uint4 v; +}; + +template +struct BytesToType {}; + +template <> +struct BytesToType<32> { + using Type = uint8; + static_assert(sizeof(Type) == 32); +}; + +template <> +struct BytesToType<16> { + using Type = uint4; + static_assert(sizeof(Type) == 16); +}; + +template <> +struct BytesToType<8> { + using Type = uint64_t; + static_assert(sizeof(Type) == 8); +}; + +template <> +struct BytesToType<4> { + using Type = uint32_t; + static_assert(sizeof(Type) == 4); +}; + +template <> +struct BytesToType<2> { + using Type = uint16_t; + static_assert(sizeof(Type) == 2); +}; + +template <> +struct BytesToType<1> { + using Type = uint8_t; + static_assert(sizeof(Type) == 1); +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct SumOp { + __device__ inline T operator()(T const& x, T const& y) { + return x + y; + } +}; + +template +struct Allreduce { + static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4); + template + static __device__ inline T run(T x, Operator& op) { + constexpr int OFFSET = THREADS / 2; + x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET)); + return Allreduce::run(x, op); + } +}; + +template <> +struct Allreduce<2> { + template + static __device__ inline T run(T x, Operator& op) { + x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1)); + return x; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// https://stackoverflow.com/questions/35311711/whats-the-right-way-to-compute-integral-base-2-logarithms-at-compile-time +constexpr int cilog2(int val) { + return val > 0 ? 1 + cilog2(val >> 1) : -1; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__device__ __forceinline__ void hadamard_mult_thread(float x[kNChunks][1 << kLogN]) { + constexpr int N = 1 << kLogN; +#pragma unroll + for (int i = 0; i < kLogN; ++i) { + const int stride = 1 << i; +#pragma unroll + for (int j = 0; j < N / 2; ++j) { + const int lo = j & (stride - 1); + const int idx = (j - lo) * 2 + lo; +#pragma unroll + for (int c = 0; c < kNChunks; ++c) { + const float a = x[c][idx]; + const float b = x[c][idx + stride]; + x[c][idx] = a + b; + x[c][idx + stride] = a - b; + } + } + } +} + +template +__device__ __forceinline__ void hadamard_mult_warp(float x[kNChunks][kNItems]) { + constexpr int N = 1 << kLogWarpSize; + int lane_id = threadIdx.x % N; +#pragma unroll + for (int step = kStepStart; step < kLogWarpSize; ++step) { + const int lane_mask = 1 << step; + const float sign = (lane_id & lane_mask) ? -1.f : 1.f; +#pragma unroll + for (int c = 0; c < kNChunks; ++c) { +#pragma unroll + for (int i = 0; i < kNItems; ++i) { + float x_val_other = __shfl_xor_sync(FULL_MASK, x[c][i], lane_mask); + x[c][i] = sign * x[c][i] + x_val_other; + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline __device__ void load_input(input_t* x, float x_vals[kNChunks][kNElts], int dim) { + using vec_t = typename BytesToType::Type; + input_t x_vals_load[kNChunks][kNElts] = {0}; +#pragma unroll + for (int c = 0; c < kNChunks; ++c) { + if ((c * blockDim.x + threadIdx.x) * kNElts < dim) { + reinterpret_cast(x_vals_load)[c] = reinterpret_cast(x)[c * blockDim.x + threadIdx.x]; + } + } +#pragma unroll + for (int c = 0; c < kNChunks; ++c) { +#pragma unroll + for (int i = 0; i < kNElts; ++i) { + x_vals[c][i] = float(x_vals_load[c][i]); + } + } +} + +template +inline __device__ void store_output(output_t* out, float out_vals[kNChunks][kNElts], int dim, float scale = 1.f) { + using vec_t = typename BytesToType::Type; + output_t out_vals_store[kNChunks][kNElts]; +#pragma unroll + for (int c = 0; c < kNChunks; ++c) { +#pragma unroll + for (int i = 0; i < kNElts; ++i) { + out_vals_store[c][i] = out_vals[c][i] * scale; + } + } +#pragma unroll + for (int c = 0; c < kNChunks; ++c) { + if ((c * blockDim.x + threadIdx.x) * kNElts < dim) { + reinterpret_cast(out)[c * blockDim.x + threadIdx.x] = reinterpret_cast(out_vals_store)[c]; + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Pre=true means the exchange before the hadamard_mult_warp, Pre=false means after. +template +inline __device__ void exchange_smem_pre(float x_vals[kNChunks][kNElts], vec_t* smem) { + constexpr int kNThreads = kWarpSize * kNWarps; + constexpr int kNExchangePerVec = kNElts / (sizeof(vec_t) / sizeof(float)); + const int warp_id = threadIdx.x / kWarpSize; + const int lane_id = threadIdx.x % kWarpSize; + const int row_t = threadIdx.x % kNWarps; + const int col_t = threadIdx.x / kNWarps; +// We use the XOR swizzle trick (new_col = col ^ row) to avoid / reduce smem bank conflicts. +#pragma unroll + for (int c0 = 0; c0 < kNChunks / kChunksPerExchange; ++c0) { + __syncthreads(); +#pragma unroll + for (int c1 = 0; c1 < kChunksPerExchange; ++c1) { +#pragma unroll + for (int r = 0; r < kNExchangePerVec; ++r) { + smem + [(c1 * kNExchangePerVec + r) * kNThreads + + (Pre ? warp_id * kWarpSize + lane_id ^ warp_id : row_t * kWarpSize + col_t ^ row_t)] = + reinterpret_cast(x_vals[c0 * kChunksPerExchange + c1])[r]; + } + } + __syncthreads(); +#pragma unroll + for (int c1 = 0; c1 < kChunksPerExchange; ++c1) { +#pragma unroll + for (int r = 0; r < kNExchangePerVec; ++r) { + reinterpret_cast(x_vals[c0 * kChunksPerExchange + c1])[r] = smem + [(c1 * kNExchangePerVec + r) * kNThreads + + (Pre ? row_t * kWarpSize + col_t ^ row_t : warp_id * kWarpSize + lane_id ^ warp_id)]; + } + } + } +} 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 new file mode 100644 index 000000000..6c3fcb0ff --- /dev/null +++ b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_cuda.cu @@ -0,0 +1,452 @@ +/****************************************************************************** + * 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); diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h new file mode 100644 index 000000000..b9f92f597 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/fast_hadamard_transform_special.h @@ -0,0 +1,298 @@ + +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ + +// Copied from https://github.com/sgl-project/fast-hadamard-transform + +// This file is auto-generated. See "code_gen.py" + +#pragma once + +__device__ __forceinline__ void hadamard_mult_thread_12(float x[12]) { + float out[12]; + out[0] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] + x[6] + x[7] + x[8] + x[9] + x[10] + x[11]; + out[1] = -x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] - x[11]; + out[2] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11]; + out[3] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11]; + out[4] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11]; + out[5] = +x[0] - x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] + x[9] - x[10] + x[11]; + out[6] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11]; + out[7] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11]; + out[8] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] + x[8] - x[9] + x[10] + x[11]; + out[9] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11]; + out[10] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11]; + out[11] = +x[0] - x[1] + x[2] - x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11]; +#pragma unroll + for (int i = 0; i < 12; i++) { + x[i] = out[i]; + } +} + +__device__ __forceinline__ void hadamard_mult_thread_20(float x[20]) { + float out[20]; + out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] + + x[14] + x[15] - x[16] + x[17] + x[18] - x[19]; + out[1] = -x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] - + x[14] - x[15] + x[16] - x[17] + x[18] + x[19]; + out[2] = -x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] - + x[14] + x[15] - x[16] + x[17] - x[18] + x[19]; + out[3] = -x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] + x[13] + + x[14] + x[15] + x[16] - x[17] + x[18] - x[19]; + out[4] = -x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + + x[14] - x[15] + x[16] + x[17] - x[18] + x[19]; + out[5] = -x[0] + x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] + + x[14] + x[15] + x[16] - x[17] - x[18] + x[19]; + out[6] = +x[0] - x[1] + x[2] + x[3] + x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - + x[14] + x[15] + x[16] + x[17] - x[18] - x[19]; + out[7] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - + x[14] - x[15] + x[16] + x[17] + x[18] - x[19]; + out[8] = +x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + + x[14] - x[15] - x[16] + x[17] + x[18] + x[19]; + out[9] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - + x[14] + x[15] - x[16] - x[17] + x[18] + x[19]; + out[10] = -x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - + x[14] - x[15] + x[16] + x[17] + x[18] + x[19]; + out[11] = -x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] - + x[14] + x[15] - x[16] + x[17] + x[18] + x[19]; + out[12] = +x[0] - x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - + x[14] + x[15] + x[16] - x[17] + x[18] + x[19]; + out[13] = +x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - + x[14] + x[15] + x[16] + x[17] - x[18] + x[19]; + out[14] = -x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + + x[14] + x[15] + x[16] + x[17] + x[18] - x[19]; + out[15] = -x[0] + x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - + x[14] + x[15] - x[16] - x[17] - x[18] - x[19]; + out[16] = +x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] - + x[14] - x[15] + x[16] - x[17] - x[18] - x[19]; + out[17] = -x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - + x[14] - x[15] - x[16] + x[17] - x[18] - x[19]; + out[18] = -x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - + x[14] - x[15] - x[16] - x[17] + x[18] - x[19]; + out[19] = +x[0] - x[1] - x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + + x[14] - x[15] - x[16] - x[17] - x[18] + x[19]; +#pragma unroll + for (int i = 0; i < 20; i++) { + x[i] = out[i]; + } +} + +__device__ __forceinline__ void hadamard_mult_thread_28(float x[28]) { + float out[28]; + out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] + + x[14] - x[15] + x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] - + x[27]; + out[1] = -x[0] + x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] - + x[14] + x[15] - x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] - + x[27]; + out[2] = -x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11] - x[12] - x[13] + + x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] + + x[27]; + out[3] = -x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] + x[11] - x[12] - x[13] - + x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] + + x[27]; + out[4] = -x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] - + x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] - + x[27]; + out[5] = -x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] + + x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] - + x[27]; + out[6] = -x[0] - x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] + x[13] - + x[14] + x[15] - x[16] - x[17] + x[18] - x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] + + x[27]; + out[7] = -x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] - x[24] - x[25] + x[26] - + x[27]; + out[8] = -x[0] - x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] + + x[27]; + out[9] = +x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] - + x[27]; + out[10] = +x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - + x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] - + x[27]; + out[11] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] - + x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] + + x[27]; + out[12] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] - x[22] - x[23] + x[24] - x[25] + x[26] - + x[27]; + out[13] = -x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] + + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] - x[26] + + x[27]; + out[14] = -x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] + + x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - + x[27]; + out[15] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] - + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] + x[26] + + x[27]; + out[16] = -x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - + x[14] - x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] - x[23] - x[24] + x[25] + x[26] + + x[27]; + out[17] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] - x[11] - x[12] + x[13] - + x[14] - x[15] - x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] - x[24] - x[25] + x[26] + + x[27]; + out[18] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] + x[11] - x[12] - x[13] - + x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] - x[25] - x[26] + + x[27]; + out[19] = -x[0] + x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] - + x[27]; + out[20] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - + x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] - + x[27]; + out[21] = -x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] - x[26] - + x[27]; + out[22] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] - x[13] + + x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] - + x[27]; + out[23] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] - + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] - x[26] - + x[27]; + out[24] = -x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] - + x[14] - x[15] + x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] - x[25] - x[26] - + x[27]; + out[25] = -x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - + x[14] - x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] - + x[27]; + out[26] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] + x[13] - + x[14] - x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] + x[26] - + x[27]; + out[27] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11] + x[12] - x[13] + + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] + + x[27]; +#pragma unroll + for (int i = 0; i < 28; i++) { + x[i] = out[i]; + } +} + +__device__ __forceinline__ void hadamard_mult_thread_40(float x[40]) { + float out[40]; + out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - + x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] - + x[27] - x[28] - x[29] - x[30] - x[31] - x[32] - x[33] - x[34] - x[35] - x[36] - x[37] - x[38] - x[39]; + out[1] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] + x[24] - x[25] - x[26] - + x[27] - x[28] + x[29] - x[30] + x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] - x[38] + x[39]; + out[2] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] + + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] + x[20] + x[21] + x[22] - x[23] + x[24] + x[25] - x[26] - + x[27] - x[28] - x[29] + x[30] - x[31] + x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] - x[39]; + out[3] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - + x[14] + x[15] + x[16] + x[17] + x[18] - x[19] + x[20] - x[21] + x[22] + x[23] - x[24] + x[25] + x[26] - + x[27] - x[28] - x[29] - x[30] + x[31] - x[32] + x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39]; + out[4] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + + x[14] - x[15] + x[16] + x[17] + x[18] + x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] + x[26] + + x[27] - x[28] - x[29] - x[30] - x[31] + x[32] - x[33] + x[34] - x[35] + x[36] + x[37] + x[38] + x[39]; + out[5] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - + x[14] + x[15] - x[16] + x[17] + x[18] + x[19] + x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] + + x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] - x[34] + x[35] - x[36] + x[37] + x[38] + x[39]; + out[6] = +x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + + x[14] - x[15] + x[16] - x[17] + x[18] + x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] + x[26] - + x[27] + x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] - x[35] + x[36] - x[37] + x[38] + x[39]; + out[7] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] - + x[14] + x[15] - x[16] + x[17] - x[18] + x[19] + x[20] + x[21] + x[22] + x[23] - x[24] - x[25] + x[26] + + x[27] - x[28] + x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] - x[36] + x[37] - x[38] + x[39]; + out[8] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - + x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] + + x[27] + x[28] - x[29] + x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] - x[37] + x[38] - x[39]; + out[9] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - + x[14] - x[15] - x[16] + x[17] - x[18] + x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] - + x[27] + x[28] + x[29] - x[30] + x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] - x[38] + x[39]; + out[10] = +x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] - + x[14] - x[15] - x[16] - x[17] + x[18] - x[19] + x[20] + x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - + x[27] - x[28] + x[29] + x[30] - x[31] + x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] - x[39]; + out[11] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] + x[22] - x[23] + x[24] + x[25] + x[26] + + x[27] - x[28] - x[29] + x[30] + x[31] - x[32] + x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39]; + out[12] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] + x[23] - x[24] + x[25] + x[26] + + x[27] + x[28] - x[29] - x[30] + x[31] + x[32] - x[33] + x[34] + x[35] - x[36] - x[37] - x[38] - x[39]; + out[13] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - + x[14] + x[15] + x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] + + x[27] + x[28] + x[29] - x[30] - x[31] + x[32] + x[33] - x[34] + x[35] + x[36] - x[37] - x[38] - x[39]; + out[14] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] + + x[27] + x[28] + x[29] + x[30] - x[31] - x[32] + x[33] + x[34] - x[35] + x[36] + x[37] - x[38] - x[39]; + out[15] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] + x[20] - x[21] - x[22] - x[23] + x[24] - x[25] + x[26] - + x[27] + x[28] + x[29] + x[30] + x[31] - x[32] - x[33] + x[34] + x[35] - x[36] + x[37] + x[38] - x[39]; + out[16] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] - + x[14] + x[15] + x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] + + x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] - x[34] + x[35] + x[36] - x[37] + x[38] + x[39]; + out[17] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] - + x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] - + x[27] + x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] - x[35] + x[36] + x[37] - x[38] + x[39]; + out[18] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] + + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] - x[25] - x[26] + + x[27] - x[28] + x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] - x[36] + x[37] + x[38] - x[39]; + out[19] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] + + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] + x[22] + x[23] - x[24] - x[25] - x[26] - + x[27] + x[28] - x[29] + x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] - x[37] + x[38] + x[39]; + out[20] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - + x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] + x[21] + x[22] + x[23] + x[24] + x[25] + x[26] + + x[27] + x[28] + x[29] + x[30] + x[31] + x[32] + x[33] + x[34] + x[35] + x[36] + x[37] + x[38] + x[39]; + out[21] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] + + x[27] + x[28] - x[29] + x[30] - x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] + x[38] - x[39]; + out[22] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] + + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] + + x[27] + x[28] + x[29] - x[30] + x[31] - x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] + x[39]; + out[23] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - + x[14] + x[15] + x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] + + x[27] + x[28] + x[29] + x[30] - x[31] + x[32] - x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39]; + out[24] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + + x[14] - x[15] + x[16] + x[17] + x[18] + x[19] - x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] - + x[27] + x[28] + x[29] + x[30] + x[31] - x[32] + x[33] - x[34] + x[35] - x[36] - x[37] - x[38] - x[39]; + out[25] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - + x[14] + x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] - + x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] + x[34] - x[35] + x[36] - x[37] - x[38] - x[39]; + out[26] = +x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + + x[14] - x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] + + x[27] - x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] + x[35] - x[36] + x[37] - x[38] - x[39]; + out[27] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] - + x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] - x[26] - + x[27] + x[28] - x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] + x[36] - x[37] + x[38] - x[39]; + out[28] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - + x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] + x[26] - + x[27] - x[28] + x[29] - x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] + x[37] - x[38] + x[39]; + out[29] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - + x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] + + x[27] - x[28] - x[29] + x[30] - x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] + x[38] - x[39]; + out[30] = +x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] - + x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] + + x[27] + x[28] - x[29] - x[30] + x[31] - x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] + x[39]; + out[31] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] + x[23] - x[24] - x[25] - x[26] - + x[27] + x[28] + x[29] - x[30] - x[31] + x[32] - x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39]; + out[32] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] - + x[27] - x[28] + x[29] + x[30] - x[31] - x[32] + x[33] - x[34] - x[35] + x[36] + x[37] + x[38] + x[39]; + out[33] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - + x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] - + x[27] - x[28] - x[29] + x[30] + x[31] - x[32] - x[33] + x[34] - x[35] - x[36] + x[37] + x[38] + x[39]; + out[34] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] + x[24] - x[25] + x[26] - + x[27] - x[28] - x[29] - x[30] + x[31] + x[32] - x[33] - x[34] + x[35] - x[36] - x[37] + x[38] + x[39]; + out[35] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] + x[25] - x[26] + + x[27] - x[28] - x[29] - x[30] - x[31] + x[32] + x[33] - x[34] - x[35] + x[36] - x[37] - x[38] + x[39]; + out[36] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] - + x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] + x[26] - + x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] + x[34] - x[35] - x[36] + x[37] - x[38] - x[39]; + out[37] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] - + x[14] - x[15] + x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] + + x[27] - x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] + x[35] - x[36] - x[37] + x[38] - x[39]; + out[38] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] + + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - + x[27] + x[28] - x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] + x[36] - x[37] - x[38] + x[39]; + out[39] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] + + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] + x[26] + + x[27] - x[28] + x[29] - x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] + x[37] - x[38] - x[39]; +#pragma unroll + for (int i = 0; i < 40; i++) { + x[i] = out[i]; + } +} diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/hadamard_jit.cuh b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/hadamard_jit.cuh new file mode 100644 index 000000000..1be821f29 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/hadamard_jit.cuh @@ -0,0 +1,482 @@ +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ + +#pragma once + +#include +#include + +#include + +#include + +#include "fast_hadamard_transform.h" +#include "fast_hadamard_transform_common.h" +#include "fast_hadamard_transform_special.h" +#include "static_switch.h" +#include +#include +#include + +namespace { + +using ::bf16_t; +using ::fp16_t; +using ::HadamardParamsBase; + +constexpr inline int ceil_log2(int val) { + int log = 0; + int p = 1; + while (p < val) { + p <<= 1; + ++log; + } + return log; +} + +template +struct FastHadamardKernelTraits { + 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; + static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t); + using vec_t = typename BytesToType::Type; + static constexpr int kNChunks = N / (kNElts * kNThreads); + static constexpr int kSmemExchangeSize = (N * 4) < (32 * 1024) ? (N * 4) : (32 * 1024); + static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; + static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); + static constexpr int kSmemSize = kSmemExchangeSize; +}; + +template +struct FastHadamardMNKernelTraits { + using input_t = input_t_; + static constexpr int kNThreads = kNThreads_; + static constexpr int kLogN = kLogN_; + static constexpr int N = (1 << kLogN) * kMultiple; + static_assert(N <= kMaxDim); + static constexpr int kNBytes = sizeof(input_t); + static_assert(kNBytes == 2 || kNBytes == 4); + static constexpr int kNElts = 4; + 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 == kMultiple); + static constexpr int kSmemExchangeSize = (N * 4) < kMaxSmem ? (N * 4) : kMaxSmem; + static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize; + static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4); + static constexpr int kSmemSize = kSmemExchangeSize; +}; + +template +using FastHadamard12NTraits = FastHadamardMNKernelTraits; + +template +using FastHadamard20NTraits = FastHadamardMNKernelTraits; + +template +using FastHadamard28NTraits = FastHadamardMNKernelTraits; + +template +using FastHadamard40NTraits = FastHadamardMNKernelTraits; + +template +SGL_DEVICE 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 +SGL_DEVICE 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 +SGL_DEVICE 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 +SGL_DEVICE 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 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 = kNThreads < 32 ? 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 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); + + extern __shared__ char smem_[]; + vec_t* smem_exchange = reinterpret_cast(smem_); + + const int batch_id = static_cast(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 +inline void set_max_dynamic_smem() { + constexpr int kSmemSize = Ktraits::kSmemSize; + if constexpr (kSmemSize >= 48 * 1024) { + auto kernel = &fast_hadamard_transform_kernel; + host::RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); + } +} + +template +inline void launch_kernel(HadamardParamsBase& params, DLDevice device) { + constexpr int kSmemSize = Ktraits::kSmemSize; + set_max_dynamic_smem(); + auto kernel = &fast_hadamard_transform_kernel; + host::LaunchKernel(dim3(params.batch), dim3(Ktraits::kNThreads), device, kSmemSize)(kernel, params); + host::RuntimeDeviceCheck(); +} + +template +inline void fast_hadamard_transform_launch(HadamardParamsBase& params, DLDevice device) { + using Ktraits = FastHadamardKernelTraits; + launch_kernel(params, device); +} + +template +inline void fast_hadamard_transform_cuda(HadamardParamsBase& params, DLDevice device) { + if (params.log_N == 3) { + fast_hadamard_transform_launch<1, 3, input_t>(params, device); + } else if (params.log_N == 4) { + fast_hadamard_transform_launch<2, 4, input_t>(params, device); + } else if (params.log_N == 5) { + fast_hadamard_transform_launch<4, 5, input_t>(params, device); + } else if (params.log_N == 6) { + fast_hadamard_transform_launch<8, 6, input_t>(params, device); + } else if (params.log_N == 7) { + fast_hadamard_transform_launch<16, 7, input_t>(params, device); + } else if (params.log_N == 8) { + fast_hadamard_transform_launch<32, 8, input_t>(params, device); + } else if (params.log_N == 9) { + fast_hadamard_transform_launch<32, 9, input_t>(params, device); + } else if (params.log_N == 10) { + fast_hadamard_transform_launch<128, 10, input_t>(params, device); + } else if (params.log_N == 11) { + fast_hadamard_transform_launch<256, 11, input_t>(params, device); + } else if (params.log_N == 12) { + fast_hadamard_transform_launch<256, 12, input_t>(params, device); + } else if (params.log_N == 13) { + fast_hadamard_transform_launch<256, 13, input_t>(params, device); + } else if (params.log_N == 14) { + fast_hadamard_transform_launch<256, 14, input_t>(params, device); + } else if (params.log_N == 15) { + fast_hadamard_transform_launch<256, 15, input_t>(params, device); + } else { + host::Panic("fast_hadamard_transform: unsupported log_N=", params.log_N); + } +} + +template +inline void fast_hadamard_transform_12N_launch(HadamardParamsBase& params, DLDevice device) { + using Ktraits = FastHadamard12NTraits; + launch_kernel(params, device); +} + +template +inline void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, DLDevice device) { + if (params.log_N == 2) { + fast_hadamard_transform_12N_launch<1, 2, input_t>(params, device); + } else if (params.log_N == 3) { + fast_hadamard_transform_12N_launch<2, 3, input_t>(params, device); + } else if (params.log_N == 4) { + fast_hadamard_transform_12N_launch<4, 4, input_t>(params, device); + } else if (params.log_N == 5) { + fast_hadamard_transform_12N_launch<8, 5, input_t>(params, device); + } else if (params.log_N == 6) { + fast_hadamard_transform_12N_launch<16, 6, input_t>(params, device); + } else if (params.log_N == 7) { + fast_hadamard_transform_12N_launch<32, 7, input_t>(params, device); + } else if (params.log_N == 8) { + fast_hadamard_transform_12N_launch<64, 8, input_t>(params, device); + } else if (params.log_N == 9) { + fast_hadamard_transform_12N_launch<128, 9, input_t>(params, device); + } else if (params.log_N == 10) { + fast_hadamard_transform_12N_launch<256, 10, input_t>(params, device); + } else { + host::Panic("fast_hadamard_transform_12N: unsupported log_N=", params.log_N); + } +} + +template +inline void fast_hadamard_transform_20N_launch(HadamardParamsBase& params, DLDevice device) { + using Ktraits = FastHadamard20NTraits; + launch_kernel(params, device); +} + +template +inline void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, DLDevice device) { + if (params.log_N == 2) { + fast_hadamard_transform_20N_launch<1, 2, input_t>(params, device); + } else if (params.log_N == 3) { + fast_hadamard_transform_20N_launch<2, 3, input_t>(params, device); + } else if (params.log_N == 4) { + fast_hadamard_transform_20N_launch<4, 4, input_t>(params, device); + } else if (params.log_N == 5) { + fast_hadamard_transform_20N_launch<8, 5, input_t>(params, device); + } else if (params.log_N == 6) { + fast_hadamard_transform_20N_launch<16, 6, input_t>(params, device); + } else if (params.log_N == 7) { + fast_hadamard_transform_20N_launch<32, 7, input_t>(params, device); + } else if (params.log_N == 8) { + fast_hadamard_transform_20N_launch<64, 8, input_t>(params, device); + } else if (params.log_N == 9) { + fast_hadamard_transform_20N_launch<128, 9, input_t>(params, device); + } else if (params.log_N == 10) { + fast_hadamard_transform_20N_launch<256, 10, input_t>(params, device); + } else { + host::Panic("fast_hadamard_transform_20N: unsupported log_N=", params.log_N); + } +} + +template +inline void fast_hadamard_transform_28N_launch(HadamardParamsBase& params, DLDevice device) { + using Ktraits = FastHadamard28NTraits; + launch_kernel(params, device); +} + +template +inline void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, DLDevice device) { + if (params.log_N == 2) { + fast_hadamard_transform_28N_launch<1, 2, input_t>(params, device); + } else if (params.log_N == 3) { + fast_hadamard_transform_28N_launch<2, 3, input_t>(params, device); + } else if (params.log_N == 4) { + fast_hadamard_transform_28N_launch<4, 4, input_t>(params, device); + } else if (params.log_N == 5) { + fast_hadamard_transform_28N_launch<8, 5, input_t>(params, device); + } else if (params.log_N == 6) { + fast_hadamard_transform_28N_launch<16, 6, input_t>(params, device); + } else if (params.log_N == 7) { + fast_hadamard_transform_28N_launch<32, 7, input_t>(params, device); + } else if (params.log_N == 8) { + fast_hadamard_transform_28N_launch<64, 8, input_t>(params, device); + } else if (params.log_N == 9) { + fast_hadamard_transform_28N_launch<128, 9, input_t>(params, device); + } else if (params.log_N == 10) { + fast_hadamard_transform_28N_launch<256, 10, input_t>(params, device); + } else { + host::Panic("fast_hadamard_transform_28N: unsupported log_N=", params.log_N); + } +} + +template +inline void fast_hadamard_transform_40N_launch(HadamardParamsBase& params, DLDevice device) { + using Ktraits = FastHadamard40NTraits; + launch_kernel(params, device); +} + +template +inline void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, DLDevice device) { + if (params.log_N == 2) { + fast_hadamard_transform_40N_launch<1, 2, input_t>(params, device); + } else if (params.log_N == 3) { + fast_hadamard_transform_40N_launch<2, 3, input_t>(params, device); + } else if (params.log_N == 4) { + fast_hadamard_transform_40N_launch<4, 4, input_t>(params, device); + } else if (params.log_N == 5) { + fast_hadamard_transform_40N_launch<8, 5, input_t>(params, device); + } else if (params.log_N == 6) { + fast_hadamard_transform_40N_launch<16, 6, input_t>(params, device); + } else if (params.log_N == 7) { + fast_hadamard_transform_40N_launch<32, 7, input_t>(params, device); + } else if (params.log_N == 8) { + fast_hadamard_transform_40N_launch<64, 8, input_t>(params, device); + } else if (params.log_N == 9) { + fast_hadamard_transform_40N_launch<128, 9, input_t>(params, device); + } else if (params.log_N == 10) { + fast_hadamard_transform_40N_launch<256, 10, input_t>(params, device); + } else { + host::Panic("fast_hadamard_transform_40N: unsupported log_N=", params.log_N); + } +} + +inline void set_hadamard_params( + HadamardParamsBase& params, + int64_t batch, + int64_t dim, + int64_t multiple, + const tvm::ffi::TensorView x, + const tvm::ffi::TensorView out, + float scale) { + std::memset(¶ms, 0, sizeof(params)); + params.batch = static_cast(batch); + params.dim = static_cast(dim); + params.log_N = ceil_log2(static_cast(dim / multiple)); + params.x_ptr = const_cast(x.data_ptr()); + params.out_ptr = const_cast(out.data_ptr()); + params.x_batch_stride = x.stride(0); + params.out_batch_stride = out.stride(0); + params.scale = scale; +} + +template +inline void run_hadamard(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { + using namespace host; + + auto N = SymbolicSize{"batch"}; + auto D = SymbolicSize{"dim"}; + auto SX = SymbolicSize{"x_batch_stride"}; + auto SO = SymbolicSize{"out_batch_stride"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, D}).with_strides({SX, 1}).with_dtype().with_device(device).verify(x); + TensorMatcher({N, D}).with_strides({SO, 1}).with_dtype().with_device(device).verify(out); + + const int64_t batch = N.unwrap(); + const int64_t dim = D.unwrap(); + + RuntimeCheck(dim % kMultiple == 0, "hadamard: dim must be divisible by ", kMultiple); + + HadamardParamsBase params; + set_hadamard_params(params, batch, dim, kMultiple, x, out, scale); + + if constexpr (kMultiple == 1) { + RuntimeCheck(dim % 8 == 0, "fast_hadamard_transform only supports hidden dim divisible by 8"); + RuntimeCheck(dim <= 32768, "fast_hadamard_transform only supports hidden dim <= 32768"); + fast_hadamard_transform_cuda(params, device.unwrap()); + } else if constexpr (kMultiple == 12) { + RuntimeCheck(dim % (4 * 12) == 0, "fast_hadamard_transform_12N only supports hidden dim divisible by 48"); + RuntimeCheck(dim <= 12 * 1024, "fast_hadamard_transform_12N only supports hidden dim <= 12288"); + fast_hadamard_transform_12N_cuda(params, device.unwrap()); + } else if constexpr (kMultiple == 20) { + RuntimeCheck(dim % (4 * 20) == 0, "fast_hadamard_transform_20N only supports hidden dim divisible by 80"); + RuntimeCheck(dim <= 20 * 1024, "fast_hadamard_transform_20N only supports hidden dim <= 20480"); + fast_hadamard_transform_20N_cuda(params, device.unwrap()); + } else if constexpr (kMultiple == 28) { + RuntimeCheck(dim % (4 * 28) == 0, "fast_hadamard_transform_28N only supports hidden dim divisible by 112"); + RuntimeCheck(dim <= 28 * 1024, "fast_hadamard_transform_28N only supports hidden dim <= 28672"); + fast_hadamard_transform_28N_cuda(params, device.unwrap()); + } else if constexpr (kMultiple == 40) { + RuntimeCheck(dim % (4 * 40) == 0, "fast_hadamard_transform_40N only supports hidden dim divisible by 160"); + RuntimeCheck(dim <= 40 * 1024, "fast_hadamard_transform_40N only supports hidden dim <= 40960"); + fast_hadamard_transform_40N_cuda(params, device.unwrap()); + } else { + Panic("Unsupported multiple"); + } +} + +template +struct HadamardKernel { + static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { + run_hadamard<1, DType>(x, out, scale); + } +}; + +template +struct Hadamard12NKernel { + static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { + run_hadamard<12, DType>(x, out, scale); + } +}; + +template +struct Hadamard20NKernel { + static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { + run_hadamard<20, DType>(x, out, scale); + } +}; + +template +struct Hadamard28NKernel { + static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { + run_hadamard<28, DType>(x, out, scale); + } +}; + +template +struct Hadamard40NKernel { + static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) { + run_hadamard<40, DType>(x, out, scale); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/csrc/fast-hadamard-transform/static_switch.h b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/static_switch.h new file mode 100644 index 000000000..aea354665 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/fast-hadamard-transform/static_switch.h @@ -0,0 +1,27 @@ +// Copied from https://github.com/sgl-project/fast-hadamard-transform + +// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h +// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h + +#pragma once + +/// @param COND - a boolean expression to switch by +/// @param CONST_NAME - a name given for the constexpr bool variable. +/// @param ... - code to execute for true and false +/// +/// Usage: +/// ``` +/// BOOL_SWITCH(flag, BoolConst, [&] { +/// some_function(...); +/// }); +/// ``` +#define BOOL_SWITCH(COND, CONST_NAME, ...) \ + [&] { \ + if (COND) { \ + static constexpr bool CONST_NAME = true; \ + return __VA_ARGS__(); \ + } else { \ + static constexpr bool CONST_NAME = false; \ + return __VA_ARGS__(); \ + } \ + }() diff --git a/python/sglang/jit_kernel/hadamard.py b/python/sglang/jit_kernel/hadamard.py new file mode 100644 index 000000000..afe9629be --- /dev/null +++ b/python/sglang/jit_kernel/hadamard.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import KERNEL_PATH, cache_once, load_jit, make_cpp_args + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_hadamard_module(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + hadamard_include_dir = (KERNEL_PATH / "csrc" / "fast-hadamard-transform").resolve() + return load_jit( + "hadamard", + *args, + cuda_files=["fast-hadamard-transform/hadamard_jit.cuh"], + cuda_wrappers=[ + ("hadamard_transform", f"HadamardKernel<{args}>::run"), + ("hadamard_transform_12n", f"Hadamard12NKernel<{args}>::run"), + ("hadamard_transform_20n", f"Hadamard20NKernel<{args}>::run"), + ("hadamard_transform_28n", f"Hadamard28NKernel<{args}>::run"), + ("hadamard_transform_40n", f"Hadamard40NKernel<{args}>::run"), + ], + extra_include_paths=[str(hadamard_include_dir)], + ) + + +def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + if not x.is_cuda: + raise RuntimeError("hadamard_transform only supports CUDA tensors") + + shapes_og = x.size() + dim_og = x.size(-1) + x = x.reshape(-1, dim_og) + if x.stride(-1) != 1: + x = x.contiguous() + + if dim_og % 8 != 0: + x = torch.nn.functional.pad(x, (0, 8 - dim_og % 8)) + dim = x.size(1) + + out = torch.empty_like(x) + module = _jit_hadamard_module(x.dtype) + module.hadamard_transform(x, out, scale) + + if dim_og % 8 != 0: + out = out[:, :dim_og] + return out.reshape(shapes_og) + + +def hadamard_transform_12n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + if not x.is_cuda: + raise RuntimeError("hadamard_transform_12n only supports CUDA tensors") + + shapes_og = x.size() + dim_og = x.size(-1) + x = x.reshape(-1, dim_og) + if x.stride(-1) != 1: + x = x.contiguous() + + pad_multiple = 4 * 12 + if dim_og % pad_multiple != 0: + x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple)) + + out = torch.empty_like(x) + module = _jit_hadamard_module(x.dtype) + module.hadamard_transform_12n(x, out, scale) + + if dim_og % pad_multiple != 0: + out = out[:, :dim_og] + return out.reshape(shapes_og) + + +def hadamard_transform_20n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + if not x.is_cuda: + raise RuntimeError("hadamard_transform_20n only supports CUDA tensors") + + shapes_og = x.size() + dim_og = x.size(-1) + x = x.reshape(-1, dim_og) + if x.stride(-1) != 1: + x = x.contiguous() + + pad_multiple = 4 * 20 + if dim_og % pad_multiple != 0: + x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple)) + + out = torch.empty_like(x) + module = _jit_hadamard_module(x.dtype) + module.hadamard_transform_20n(x, out, scale) + + if dim_og % pad_multiple != 0: + out = out[:, :dim_og] + return out.reshape(shapes_og) + + +def hadamard_transform_28n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + if not x.is_cuda: + raise RuntimeError("hadamard_transform_28n only supports CUDA tensors") + + shapes_og = x.size() + dim_og = x.size(-1) + x = x.reshape(-1, dim_og) + if x.stride(-1) != 1: + x = x.contiguous() + + pad_multiple = 4 * 28 + if dim_og % pad_multiple != 0: + x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple)) + + out = torch.empty_like(x) + module = _jit_hadamard_module(x.dtype) + module.hadamard_transform_28n(x, out, scale) + + if dim_og % pad_multiple != 0: + out = out[:, :dim_og] + return out.reshape(shapes_og) + + +def hadamard_transform_40n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + if not x.is_cuda: + raise RuntimeError("hadamard_transform_40n only supports CUDA tensors") + + shapes_og = x.size() + dim_og = x.size(-1) + x = x.reshape(-1, dim_og) + if x.stride(-1) != 1: + x = x.contiguous() + + pad_multiple = 4 * 40 + if dim_og % pad_multiple != 0: + x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple)) + + out = torch.empty_like(x) + module = _jit_hadamard_module(x.dtype) + module.hadamard_transform_40n(x, out, scale) + + if dim_og % pad_multiple != 0: + out = out[:, :dim_og] + return out.reshape(shapes_og) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 8b48dfd91..ca54a931b 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -125,7 +125,7 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: if _is_hip: from fast_hadamard_transform import hadamard_transform else: - from sgl_kernel import hadamard_transform + from sglang.jit_kernel.hadamard import hadamard_transform hidden_size = x.size(-1) assert ( diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 0c8d17802..4468c8524 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -106,15 +106,6 @@ FetchContent_Declare( ) FetchContent_Populate(repo-mscclpp) -# fast-hadamard-transform -FetchContent_Declare( - repo-fast-hadamard-transform - GIT_REPOSITORY https://github.com/sgl-project/fast-hadamard-transform.git - GIT_TAG 48f3c13764dc2ec662ade842a4696a90a137f1bc - GIT_SHALLOW OFF -) -FetchContent_Populate(repo-fast-hadamard-transform) - # ccache option option(ENABLE_CCACHE "Whether to use ccache" ON) find_program(CCACHE_FOUND ccache) @@ -343,9 +334,6 @@ set(SOURCES "${repo-flashinfer_SOURCE_DIR}/csrc/renorm.cu" "${repo-flashinfer_SOURCE_DIR}/csrc/sampling.cu" - "${repo-fast-hadamard-transform_SOURCE_DIR}/csrc/fast_hadamard_transform_cuda.cu" - "${repo-fast-hadamard-transform_SOURCE_DIR}/csrc/fast_hadamard_transform.cpp" - "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_causal_sm80.cu" "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_sm80.cu" "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_fp16_causal_sm80.cu" diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index d0b6fcf80..004b85304 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -591,24 +591,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "es_sm100_mxfp8_blockscaled_grouped_quant(Tensor input, Tensor problem_sizes, Tensor expert_offsets, Tensor " "blockscale_offsets, Tensor quant_output, Tensor scale_factor) -> () "); m.impl("es_sm100_mxfp8_blockscaled_grouped_quant", &es_sm100_mxfp8_blockscaled_grouped_quant); - - /* - * From fast-hadamard-transform - */ - m.def("fast_hadamard_transform(Tensor x, float scale) -> Tensor"); - m.impl("fast_hadamard_transform", torch::kCUDA, &fast_hadamard_transform); - - m.def("fast_hadamard_transform_12N(Tensor x, float scale) -> Tensor"); - m.impl("fast_hadamard_transform_12N", torch::kCUDA, &fast_hadamard_transform_12N); - - m.def("fast_hadamard_transform_20N(Tensor x, float scale) -> Tensor"); - m.impl("fast_hadamard_transform_20N", torch::kCUDA, &fast_hadamard_transform_20N); - - m.def("fast_hadamard_transform_28N(Tensor x, float scale) -> Tensor"); - m.impl("fast_hadamard_transform_28N", torch::kCUDA, &fast_hadamard_transform_28N); - - m.def("fast_hadamard_transform_40N(Tensor x, float scale) -> Tensor"); - m.impl("fast_hadamard_transform_40N", torch::kCUDA, &fast_hadamard_transform_40N); } REGISTER_EXTENSION(common_ops) diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 2eb0856aa..3b10d90a9 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -936,15 +936,6 @@ void es_sm100_mxfp8_blockscaled_grouped_quant( torch::Tensor& quant_output, torch::Tensor& scale_factor); -/* - * From fast-hadamard-transform - */ -torch::Tensor fast_hadamard_transform(torch::Tensor& x, double scale); -torch::Tensor fast_hadamard_transform_12N(torch::Tensor& x, double scale); -torch::Tensor fast_hadamard_transform_20N(torch::Tensor& x, double scale); -torch::Tensor fast_hadamard_transform_28N(torch::Tensor& x, double scale); -torch::Tensor fast_hadamard_transform_40N(torch::Tensor& x, double scale); - /* * From flashmla */ diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index a24d3573b..dc97be8c0 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -65,13 +65,6 @@ from sgl_kernel.gemm import ( silu_and_mul_scaled_fp4_grouped_quant, ) from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda -from sgl_kernel.hadamard import ( - hadamard_transform, - hadamard_transform_12n, - hadamard_transform_20n, - hadamard_transform_28n, - hadamard_transform_40n, -) from sgl_kernel.kvcacheio import ( transfer_kv_all_layer, transfer_kv_all_layer_mla, diff --git a/sgl-kernel/python/sgl_kernel/hadamard.py b/sgl-kernel/python/sgl_kernel/hadamard.py deleted file mode 100644 index 102c540f9..000000000 --- a/sgl-kernel/python/sgl_kernel/hadamard.py +++ /dev/null @@ -1,21 +0,0 @@ -import torch - - -def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - return torch.ops.sgl_kernel.fast_hadamard_transform.default(x, scale) - - -def hadamard_transform_12n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - return torch.ops.sgl_kernel.fast_hadamard_transform_12N.default(x, scale) - - -def hadamard_transform_20n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - return torch.ops.sgl_kernel.fast_hadamard_transform_20N.default(x, scale) - - -def hadamard_transform_28n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - return torch.ops.sgl_kernel.fast_hadamard_transform_28N.default(x, scale) - - -def hadamard_transform_40n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - return torch.ops.sgl_kernel.fast_hadamard_transform_40N.default(x, scale) diff --git a/sgl-kernel/tests/test_hadamard.py b/sgl-kernel/tests/test_hadamard.py index 5d1cd40e2..a70e8f1f1 100644 --- a/sgl-kernel/tests/test_hadamard.py +++ b/sgl-kernel/tests/test_hadamard.py @@ -5,7 +5,14 @@ import torch import torch.nn.functional as F from einops import rearrange, repeat from scipy.linalg import hadamard -from sgl_kernel import hadamard_transform + +try: + from sgl_kernel import hadamard_transform +except Exception: + pytest.skip( + "sgl-kernel hadamard interface was removed (migrated to jit_kernel)", + allow_module_level=True, + ) def hadamard_transform_ref(x, scale=1.0): diff --git a/test/registered/kernels/test_nsa_indexer.py b/test/registered/kernels/test_nsa_indexer.py index 07504a684..0a1e59654 100644 --- a/test/registered/kernels/test_nsa_indexer.py +++ b/test/registered/kernels/test_nsa_indexer.py @@ -584,8 +584,8 @@ class TestNSAIndexer(CustomTestCase): output = rotate_activation(x) self.assertEqual(output.shape, x.shape) self.assertEqual(output.dtype, torch.bfloat16) - except ImportError: - self.skipTest("sgl_kernel not available for hadamard_transform") + except Exception: + self.skipTest("hadamard JIT kernel not available") def test_rotate_activation_invalid_size(self): """Test that rotate_activation fails with non-power-of-2 size."""