chore: import upstream snapshot with attribution

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()

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;
+};

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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#define FULL_MASK 0xffffffff
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct uint8 {
+  uint4 u;
+  uint4 v;
+};
+
+template <int BYTES>
+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 <typename T>
+struct SumOp {
+  __device__ inline T operator()(T const& x, T const& y) {
+    return x + y;
+  }
+};
+
+template <int THREADS>
+struct Allreduce {
+  static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
+  template <typename T, typename Operator>
+  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<OFFSET>::run(x, op);
+  }
+};
+
+template <>
+struct Allreduce<2> {
+  template <typename T, typename Operator>
+  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 <int kLogN, int kNChunks>
+__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 <int kLogWarpSize, int kStepStart, int kNChunks, int kNItems>
+__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 <int kNChunks, int kNElts, typename input_t>
+inline __device__ void load_input(input_t* x, float x_vals[kNChunks][kNElts], int dim) {
+  using vec_t = typename BytesToType<sizeof(input_t) * kNElts>::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<vec_t*>(x_vals_load)[c] = reinterpret_cast<const vec_t*>(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 <int kNChunks, int kNElts, typename output_t>
+inline __device__ void store_output(output_t* out, float out_vals[kNChunks][kNElts], int dim, float scale = 1.f) {
+  using vec_t = typename BytesToType<sizeof(output_t) * kNElts>::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<vec_t*>(out)[c * blockDim.x + threadIdx.x] = reinterpret_cast<const vec_t*>(out_vals_store)[c];
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Pre=true means the exchange before the hadamard_mult_warp, Pre=false means after.
+template <int kNChunks, int kChunksPerExchange, int kNElts, int kWarpSize, int kNWarps, bool Pre, typename vec_t>
+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<vec_t*>(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<vec_t*>(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)];
+      }
+    }
+  }
+}

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];
+  }
+}

python/sglang/jit_kernel/csrc/fast-hadamard-transform/hadamard_jit.cuh

@@ -0,0 +1,482 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include "fast_hadamard_transform.h"
+#include "fast_hadamard_transform_common.h"
+#include "fast_hadamard_transform_special.h"
+#include "static_switch.h"
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+
+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 <int kNThreads_, int kLogN_, typename input_t_>
+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<kNBytes * kNElts>::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 <int kNThreads_, int kLogN_, int kMultiple, int kMaxDim, int kMaxSmem, typename input_t_>
+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<kNBytes * kNElts>::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 <int kNThreads_, int kLogN_, typename input_t_>
+using FastHadamard12NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 12, 12 * 1024, 24 * 1024, input_t_>;
+
+template <int kNThreads_, int kLogN_, typename input_t_>
+using FastHadamard20NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 20, 20 * 1024, 40 * 1024, input_t_>;
+
+template <int kNThreads_, int kLogN_, typename input_t_>
+using FastHadamard28NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 28, 28 * 1024, 28 * 1024, input_t_>;
+
+template <int kNThreads_, int kLogN_, typename input_t_>
+using FastHadamard40NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 40, 40 * 1024, 40 * 1024, input_t_>;
+
+template <int kNChunks>
+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 <int kNChunks>
+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 <int kNChunks>
+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 <int kNChunks>
+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 <typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads) void fast_hadamard_transform_kernel(HadamardParamsBase params) {
+  constexpr int kNThreads = Ktraits::kNThreads;
+  constexpr int kNElts = Ktraits::kNElts;
+  constexpr int kNExchangePerVec = Ktraits::kNExchangePerVec;
+  constexpr int 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<vec_t*>(smem_);
+
+  const int batch_id = static_cast<int>(blockIdx.x);
+  input_t* x = reinterpret_cast<input_t*>(params.x_ptr) + batch_id * params.x_batch_stride;
+  input_t* out = reinterpret_cast<input_t*>(params.out_ptr) + batch_id * params.out_batch_stride;
+
+  float x_vals[kNChunks][kNElts];
+  load_input<kNChunks, kNElts, input_t>(x, x_vals, params.dim);
+
+  hadamard_mult_thread<kLogNElts, kNChunks>(x_vals);
+  hadamard_mult_warp<kLogWarpSize, 0, kNChunks, kNElts>(x_vals);
+
+  if constexpr (kNWarps > 1) {
+    exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, true, vec_t>(x_vals, smem_exchange);
+    hadamard_mult_warp<kLogNWarps, 0, kNChunks, kNElts>(x_vals);
+    exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, false, vec_t>(x_vals, smem_exchange);
+  }
+
+  if constexpr (kNChunks > 1) {
+    float x_vals_transposed[kNElts][kNChunks];
+#pragma unroll
+    for (int c = 0; c < kNChunks; ++c) {
+#pragma unroll
+      for (int i = 0; i < kNElts; ++i) {
+        x_vals_transposed[i][c] = x_vals[c][i];
+      }
+    }
+
+    if constexpr (kNChunks == 12) {
+      hadamard_mult_thread_chunk_12<kNElts>(x_vals_transposed);
+    } else if constexpr (kNChunks == 20) {
+      hadamard_mult_thread_chunk_20<kNElts>(x_vals_transposed);
+    } else if constexpr (kNChunks == 28) {
+      hadamard_mult_thread_chunk_28<kNElts>(x_vals_transposed);
+    } else if constexpr (kNChunks == 40) {
+      hadamard_mult_thread_chunk_40<kNElts>(x_vals_transposed);
+    } else {
+      constexpr int kLogNChunks = cilog2(kNChunks);
+      static_assert(1 << kLogNChunks == kNChunks, "kNChunks must be a power of 2");
+      hadamard_mult_thread<kLogNChunks, kNElts>(x_vals_transposed);
+    }
+
+#pragma unroll
+    for (int c = 0; c < kNChunks; ++c) {
+#pragma unroll
+      for (int i = 0; i < kNElts; ++i) {
+        x_vals[c][i] = x_vals_transposed[i][c];
+      }
+    }
+  }
+
+  store_output<kNChunks, kNElts, input_t>(out, x_vals, params.dim, params.scale);
+}
+
+template <typename Ktraits>
+inline void set_max_dynamic_smem() {
+  constexpr int kSmemSize = Ktraits::kSmemSize;
+  if constexpr (kSmemSize >= 48 * 1024) {
+    auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
+    host::RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+  }
+}
+
+template <typename Ktraits>
+inline void launch_kernel(HadamardParamsBase& params, DLDevice device) {
+  constexpr int kSmemSize = Ktraits::kSmemSize;
+  set_max_dynamic_smem<Ktraits>();
+  auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
+  host::LaunchKernel(dim3(params.batch), dim3(Ktraits::kNThreads), device, kSmemSize)(kernel, params);
+  host::RuntimeDeviceCheck();
+}
+
+template <int kNThreads, int kLogN, typename input_t>
+inline void fast_hadamard_transform_launch(HadamardParamsBase& params, DLDevice device) {
+  using Ktraits = FastHadamardKernelTraits<kNThreads, kLogN, input_t>;
+  launch_kernel<Ktraits>(params, device);
+}
+
+template <typename input_t>
+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 <int kNThreads, int kLogN, typename input_t>
+inline void fast_hadamard_transform_12N_launch(HadamardParamsBase& params, DLDevice device) {
+  using Ktraits = FastHadamard12NTraits<kNThreads, kLogN, input_t>;
+  launch_kernel<Ktraits>(params, device);
+}
+
+template <typename input_t>
+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 <int kNThreads, int kLogN, typename input_t>
+inline void fast_hadamard_transform_20N_launch(HadamardParamsBase& params, DLDevice device) {
+  using Ktraits = FastHadamard20NTraits<kNThreads, kLogN, input_t>;
+  launch_kernel<Ktraits>(params, device);
+}
+
+template <typename input_t>
+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 <int kNThreads, int kLogN, typename input_t>
+inline void fast_hadamard_transform_28N_launch(HadamardParamsBase& params, DLDevice device) {
+  using Ktraits = FastHadamard28NTraits<kNThreads, kLogN, input_t>;
+  launch_kernel<Ktraits>(params, device);
+}
+
+template <typename input_t>
+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 <int kNThreads, int kLogN, typename input_t>
+inline void fast_hadamard_transform_40N_launch(HadamardParamsBase& params, DLDevice device) {
+  using Ktraits = FastHadamard40NTraits<kNThreads, kLogN, input_t>;
+  launch_kernel<Ktraits>(params, device);
+}
+
+template <typename input_t>
+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(&params, 0, sizeof(params));
+  params.batch = static_cast<int>(batch);
+  params.dim = static_cast<int>(dim);
+  params.log_N = ceil_log2(static_cast<int>(dim / multiple));
+  params.x_ptr = const_cast<void*>(x.data_ptr());
+  params.out_ptr = const_cast<void*>(out.data_ptr());
+  params.x_batch_stride = x.stride(0);
+  params.out_batch_stride = out.stride(0);
+  params.scale = scale;
+}
+
+template <int kMultiple, typename DType>
+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<kDLCUDA>();
+
+  TensorMatcher({N, D}).with_strides({SX, 1}).with_dtype<DType>().with_device(device).verify(x);
+  TensorMatcher({N, D}).with_strides({SO, 1}).with_dtype<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<DType>(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<DType>(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<DType>(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<DType>(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<DType>(params, device.unwrap());
+  } else {
+    Panic("Unsupported multiple");
+  }
+}
+
+template <typename DType>
+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 <typename DType>
+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 <typename DType>
+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 <typename DType>
+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 <typename DType>
+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

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<BoolConst>(...);
+/// });
+/// ```
+#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__();                     \
+    }                                           \
+  }()