Remove source

README.md

@@ -1,10 +1,13 @@
 ---
 tags:
-- kernel
+  - kernel
 ---
 
 ![Status](https://hubwebhook.dholtz.com/shield?repo=kernels-community/activation)
 
 ## Activation
 
-Activation kernels from [vLLM](https://github.com/vllm-project/vllm/blob/main/csrc/activation_kernels.cu).
+Activation kernels from [vLLM](https://github.com/vllm-project/vllm/blob/main/csrc/activation_kernels.cu).
+
+Kernel source: https://github.com/huggingface/kernels-community/tree/main/activation
+

activation/activation_kernels.cu

@@ -1,244 +0,0 @@
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
-
-#include <cmath>
-
-#include "cuda_compat.h"
-#include "dispatch_utils.h"
-
-namespace vllm {
-
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first>
-__device__ __forceinline__ scalar_t compute(const scalar_t& x,
-                                            const scalar_t& y) {
-  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
-}
-// Activation and gating kernel template.
-
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
-          bool act_first>
-__global__ void act_and_mul_kernel(
-    scalar_t* __restrict__ out,          // [..., d]
-    const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d) {
-  const int64_t token_idx = blockIdx.x;
-  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
-    const scalar_t x = VLLM_LDG(&input[token_idx * 2 * d + idx]);
-    const scalar_t y = VLLM_LDG(&input[token_idx * 2 * d + d + idx]);
-    out[token_idx * d + idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
-  }
-}
-
-template <typename T>
-__device__ __forceinline__ T silu_kernel(const T& x) {
-  // x * sigmoid(x)
-  return (T)(((float)x) / (1.0f + expf((float)-x)));
-}
-
-template <typename T>
-__device__ __forceinline__ T gelu_kernel(const T& x) {
-  // Equivalent to PyTorch GELU with 'none' approximation.
-  // Refer to:
-  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
-  const float f = (float)x;
-  constexpr float ALPHA = M_SQRT1_2;
-  return (T)(f * 0.5f * (1.0f + erf(f * ALPHA)));
-}
-
-template <typename T>
-__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
-  // Equivalent to PyTorch GELU with 'tanh' approximation.
-  // Refer to:
-  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
-  const float f = (float)x;
-  constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
-  constexpr float KAPPA = 0.044715;
-  float x_cube = f * f * f;
-  float inner = BETA * (f + KAPPA * x_cube);
-  return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
-}
-
-}  // namespace vllm
-
-// Launch activation and gating kernel.
-// Use ACT_FIRST (bool) indicating whether to apply the activation function
-// first.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST)                 \
-  int d = input.size(-1) / 2;                                            \
-  int64_t num_tokens = input.numel() / input.size(-1);                   \
-  dim3 grid(num_tokens);                                                 \
-  dim3 block(std::min(d, 1024));                                         \
-  if (num_tokens == 0) {                                                 \
-    return;                                                              \
-  }                                                                      \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
-  VLLM_DISPATCH_FLOATING_TYPES(                                          \
-      input.scalar_type(), "act_and_mul_kernel", [&] {                   \
-        vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST>  \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
-                                         input.data_ptr<scalar_t>(), d); \
-      });
-
-void silu_and_mul(torch::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
-{
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, true);
-}
-
-void mul_and_silu(torch::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
-{
-  // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
-  // applies the silu to the latter half of the input.
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, false);
-}
-
-void gelu_and_mul(torch::Tensor& out,    // [..., d]
-                  torch::Tensor& input)  // [..., 2 * d]
-{
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, true);
-}
-
-void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
-                       torch::Tensor& input)  // [..., 2 * d]
-{
-  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel, true);
-}
-
-namespace vllm {
-
-template <typename T>
-__device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) {
-  const float f = (float)x;
-  return (T)(f > threshold ? f : 0.0f);
-}
-
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float)>
-__global__ void act_and_mul_kernel_with_param(
-    scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
-    const float param) {
-  const int64_t token_idx = blockIdx.x;
-  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
-    const scalar_t x = VLLM_LDG(&input[token_idx * 2 * d + idx]);
-    const scalar_t y = VLLM_LDG(&input[token_idx * 2 * d + d + idx]);
-    out[token_idx * d + idx] = ACT_FN(x, param) * y;
-  }
-}
-
-}  // namespace vllm
-
-#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PARAM)         \
-  int d = input.size(-1) / 2;                                           \
-  int64_t num_tokens = input.numel() / input.size(-1);                  \
-  dim3 grid(num_tokens);                                                \
-  dim3 block(std::min(d, 1024));                                        \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));     \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();         \
-  VLLM_DISPATCH_FLOATING_TYPES(                                         \
-      input.scalar_type(), "act_and_mul_kernel_with_param", [&] {       \
-        vllm::act_and_mul_kernel_with_param<scalar_t, KERNEL<scalar_t>> \
-            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),      \
-                                         input.data_ptr<scalar_t>(), d, \
-                                         PARAM);                        \
-      });
-
-void fatrelu_and_mul(torch::Tensor& out,    // [..., d],
-                     torch::Tensor& input,  // [..., 2 * d]
-                     double threshold) {
-  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(vllm::fatrelu_kernel, threshold);
-}
-namespace vllm {
-
-// Element-wise activation kernel template.
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&)>
-__global__ void activation_kernel(
-    scalar_t* __restrict__ out,          // [..., d]
-    const scalar_t* __restrict__ input,  // [..., d]
-    const int d) {
-  const int64_t token_idx = blockIdx.x;
-  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
-    const scalar_t x = VLLM_LDG(&input[token_idx * d + idx]);
-    out[token_idx * d + idx] = ACT_FN(x);
-  }
-}
-
-}  // namespace vllm
-
-// Launch element-wise activation kernel.
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                       \
-  int d = input.size(-1);                                                      \
-  int64_t num_tokens = input.numel() / d;                                      \
-  dim3 grid(num_tokens);                                                       \
-  dim3 block(std::min(d, 1024));                                               \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
-  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \
-    vllm::activation_kernel<scalar_t, KERNEL<scalar_t>>                        \
-        <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),                 \
-                                     input.data_ptr<scalar_t>(), d);           \
-  });
-
-namespace vllm {
-
-
-template <typename T>
-__device__ __forceinline__ T gelu_new_kernel(const T& x) {
-  const float x3 = (float)(x * x * x);
-  const T t = (T)tanhf((T)(0.79788456f * (float)(x + (T)(0.044715f * x3))));
-  return ((T)0.5) * x * (((T)1.0) + t);
-}
-
-template <typename T>
-__device__ __forceinline__ T gelu_fast_kernel(const T& x) {
-  const float f = (float)x;
-  const T t =
-      (T)tanhf(((T)(f * 0.79788456f)) * (((T)1.0) + (T)(0.044715f * f) * x));
-  return ((T)0.5) * x * (((T)1.0) + t);
-}
-
-template <typename T>
-__device__ __forceinline__ T gelu_quick_kernel(const T& x) {
-  // x * sigmoid(1.702 * x)
-  return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x)));
-}
-
-}  // namespace vllm
-
-void gelu_new(torch::Tensor& out,    // [..., d]
-              torch::Tensor& input)  // [..., d]
-{
-  LAUNCH_ACTIVATION_KERNEL(vllm::gelu_new_kernel);
-}
-
-void gelu_fast(torch::Tensor& out,    // [..., d]
-               torch::Tensor& input)  // [..., d]
-{
-  LAUNCH_ACTIVATION_KERNEL(vllm::gelu_fast_kernel);
-}
-
-void gelu_quick(torch::Tensor& out,    // [..., d]
-                torch::Tensor& input)  // [..., d]
-{
-  LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
-}
-
-void gelu(torch::Tensor& out,    // [..., d]
-          torch::Tensor& input)  // [..., d]
-{
-  LAUNCH_ACTIVATION_KERNEL(vllm::gelu_kernel);
-}
-
-void gelu_tanh(torch::Tensor& out,    // [..., d]
-               torch::Tensor& input)  // [..., d]
-{
-  LAUNCH_ACTIVATION_KERNEL(vllm::gelu_tanh_kernel);
-}
-
-void silu(torch::Tensor& out,    // [..., d]
-          torch::Tensor& input)  // [..., d]
-{
-  LAUNCH_ACTIVATION_KERNEL(vllm::silu_kernel);
-}

activation/cuda_compat.h

@@ -1,49 +0,0 @@
-#pragma once
-
-#ifdef USE_ROCM
-  #include <hip/hip_runtime.h>
-#endif
-
-#if defined(USE_ROCM) && defined(__GFX9__)
-  #define WARP_SIZE 64
-#else
-  #define WARP_SIZE 32
-#endif
-
-#ifndef USE_ROCM
-  #define VLLM_LDG(arg) __ldg(arg)
-#else
-  #define VLLM_LDG(arg) *(arg)
-#endif
-
-#ifndef USE_ROCM
-  #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \
-    __shfl_xor_sync(uint32_t(-1), var, lane_mask)
-  #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \
-    __shfl_xor_sync(uint32_t(-1), var, lane_mask, width)
-#else
-  #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask)
-  #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \
-    __shfl_xor(var, lane_mask, width)
-#endif
-
-#ifndef USE_ROCM
-  #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane)
-#else
-  #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane)
-#endif
-
-#ifndef USE_ROCM
-  #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \
-    __shfl_down_sync(uint32_t(-1), var, lane_delta)
-#else
-  #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) __shfl_down(var, lane_delta)
-#endif
-
-#ifndef USE_ROCM
-  #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \
-    cudaFuncSetAttribute(FUNC, cudaFuncAttributeMaxDynamicSharedMemorySize, VAL)
-#else
-  #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \
-    hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL)
-#endif

activation/dispatch_utils.h

@@ -1,83 +0,0 @@
-/*
- * Adapted from
- * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h
- */
-#pragma once
-
-#include <torch/all.h>
-
-// Need a special dispatch case macro since we will nest the FP8 dispatch.
-// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'.
-#define AT_DISPATCH_FP8_CASE(enum_type, ...) \
-  AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
-
-#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
-
-#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
-
-// ROCm devices might use either fn or fnuz, so set up dispatch table for both.
-// A host-based check at runtime will create a preferred FP8 type for ROCm
-// such that the correct kernel is dispatched.
-#ifdef USE_ROCM
-  #define VLLM_DISPATCH_CASE_FP8_TYPES(...)                          \
-    AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
-    AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__)
-
-  #define VLLM_DISPATCH_CASE_QUANT_TYPES(...)                      \
-    AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__)   \
-    AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \
-    AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
-#else
-  #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
-    AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__)
-
-  #define VLLM_DISPATCH_CASE_QUANT_TYPES(...)                    \
-    AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
-    AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
-#endif
-
-// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'.
-// See AT_DISPATCH_FP8_CASE above.
-#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
-
-#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
-
-#define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...)   \
-  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
-  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)     \
-  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__)
-
-#define VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME,                               \
-                     VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(__VA_ARGS__))
-
-#define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...)         \
-  AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__)  \
-  AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)  \
-  AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \
-  AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__)   \
-  AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__)
-
-#define VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \
-  AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__)       \
-  AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)       \
-  AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__)      \
-  AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__)        \
-  AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__)       \
-  AT_DISPATCH_CASE(at::ScalarType::UInt16, __VA_ARGS__)     \
-  AT_DISPATCH_CASE(at::ScalarType::UInt32, __VA_ARGS__)     \
-  AT_DISPATCH_CASE(at::ScalarType::UInt64, __VA_ARGS__)
-
-#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
-
-#define VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \
-  AT_DISPATCH_SWITCH(                                              \
-      TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__))

build.toml

@@ -1,18 +0,0 @@
-[general]
-name = "activation"
-universal = false
-
-[torch]
-src = [
-    "torch-ext/torch_binding.cpp",
-    "torch-ext/torch_binding.h",
-]
-
-[kernel.activation]
-backend = "cuda"
-depends = ["torch"]
-src = [
-    "activation/activation_kernels.cu",
-    "activation/cuda_compat.h",
-    "activation/dispatch_utils.h",
-]

flake.lock

@@ -1,168 +0,0 @@
-{
-  "nodes": {
-    "flake-compat": {
-      "locked": {
-        "lastModified": 1747046372,
-        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
-        "type": "github"
-      },
-      "original": {
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "type": "github"
-      }
-    },
-    "flake-compat_2": {
-      "locked": {
-        "lastModified": 1747046372,
-        "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=",
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885",
-        "type": "github"
-      },
-      "original": {
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "type": "github"
-      }
-    },
-    "flake-utils": {
-      "inputs": {
-        "systems": "systems"
-      },
-      "locked": {
-        "lastModified": 1731533236,
-        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
-        "type": "github"
-      },
-      "original": {
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "type": "github"
-      }
-    },
-    "flake-utils_2": {
-      "inputs": {
-        "systems": "systems_2"
-      },
-      "locked": {
-        "lastModified": 1731533236,
-        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
-        "type": "github"
-      },
-      "original": {
-        "owner": "numtide",
-        "repo": "flake-utils",
-        "type": "github"
-      }
-    },
-    "hf-nix": {
-      "inputs": {
-        "flake-compat": "flake-compat_2",
-        "flake-utils": "flake-utils_2",
-        "nixpkgs": "nixpkgs"
-      },
-      "locked": {
-        "lastModified": 1759493343,
-        "narHash": "sha256-8fhl0gwMAnOkQbogPIVq+Fha+Yeq52FaRXfwF+F9Q+k=",
-        "owner": "huggingface",
-        "repo": "hf-nix",
-        "rev": "b1fc3a18b52447a0f24bc6884418edc5e66082b9",
-        "type": "github"
-      },
-      "original": {
-        "owner": "huggingface",
-        "repo": "hf-nix",
-        "type": "github"
-      }
-    },
-    "kernel-builder": {
-      "inputs": {
-        "flake-compat": "flake-compat",
-        "flake-utils": "flake-utils",
-        "hf-nix": "hf-nix",
-        "nixpkgs": [
-          "kernel-builder",
-          "hf-nix",
-          "nixpkgs"
-        ]
-      },
-      "locked": {
-        "lastModified": 1759516823,
-        "narHash": "sha256-UJVvZHtS9c64Dm4iZRaOKWB+VHI7jzcazGH57KXWeg8=",
-        "owner": "huggingface",
-        "repo": "kernel-builder",
-        "rev": "e13610a05f67b7296be9ead89ad172a0a088a1c3",
-        "type": "github"
-      },
-      "original": {
-        "owner": "huggingface",
-        "repo": "kernel-builder",
-        "type": "github"
-      }
-    },
-    "nixpkgs": {
-      "locked": {
-        "lastModified": 1755963616,
-        "narHash": "sha256-6yD0ww/S8n+U2uPYcJZ3DRURP8Kx036GRpR2uPNZroE=",
-        "owner": "nixos",
-        "repo": "nixpkgs",
-        "rev": "73e96df7cff5783f45e21342a75a1540c4eddce4",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nixos",
-        "ref": "nixos-unstable-small",
-        "repo": "nixpkgs",
-        "type": "github"
-      }
-    },
-    "root": {
-      "inputs": {
-        "kernel-builder": "kernel-builder"
-      }
-    },
-    "systems": {
-      "locked": {
-        "lastModified": 1681028828,
-        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
-        "owner": "nix-systems",
-        "repo": "default",
-        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-systems",
-        "repo": "default",
-        "type": "github"
-      }
-    },
-    "systems_2": {
-      "locked": {
-        "lastModified": 1681028828,
-        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
-        "owner": "nix-systems",
-        "repo": "default",
-        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-systems",
-        "repo": "default",
-        "type": "github"
-      }
-    }
-  },
-  "root": "root",
-  "version": 7
-}

flake.nix

@@ -1,17 +0,0 @@
-{
-  description = "Flake for activation kernels";
-
-  inputs = {
-    kernel-builder.url = "github:huggingface/kernel-builder";
-  };
-
-  outputs =
-    {
-      self,
-      kernel-builder,
-    }:
-    kernel-builder.lib.genFlakeOutputs {
-      inherit self;
-      path = ./.;
-    };
-}

tests/kernels/allclose_default.py

@@ -1,14 +0,0 @@
-import torch
-
-# Reference default values of atol and rtol are from
-# https://github.com/pytorch/pytorch/blob/6d96beb6bec24d73ee3f080bac54d2104068f675/test/test_transformers.py#L67
-default_atol = {torch.float16: 1e-3, torch.bfloat16: 1e-3, torch.float: 1e-5}
-default_rtol = {torch.float16: 1e-3, torch.bfloat16: 1.6e-2, torch.float: 1.3e-6}
-
-
-def get_default_atol(output) -> float:
-    return default_atol[output.dtype]
-
-
-def get_default_rtol(output) -> float:
-    return default_rtol[output.dtype]

tests/kernels/test_activation.py

@@ -1,206 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import math
-import random
-from typing import Type
-
-import activation
-import pytest
-import torch
-import torch.nn.functional as F
-
-from .utils import opcheck
-from .allclose_default import get_default_atol, get_default_rtol
-
-DTYPES = [torch.half, torch.bfloat16, torch.float]
-NUM_TOKENS = [7, 83, 2048]  # Arbitrary values for testing
-D = [512, 13824]  # Arbitrary values for testing
-SEEDS = [0]
-CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
-
-
-def gelu_fast(x: torch.Tensor) -> torch.Tensor:
-    return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))
-
-
-def gelu_new(x: torch.Tensor) -> torch.Tensor:
-    c = math.sqrt(2.0 / math.pi)
-    return 0.5 * x * (1.0 + torch.tanh(c * (x + 0.044715 * torch.pow(x, 3.0))))
-
-
-def gelu_quick(x: torch.Tensor) -> torch.Tensor:
-    return x * torch.sigmoid(1.702 * x)
-
-
-def fatrelu_and_mul(x: torch.Tensor, threshold: float) -> torch.Tensor:
-    d = x.shape[-1] // 2
-    x1 = x[..., :d]
-    x2 = x[..., d:]
-    x1 = F.threshold(x1, threshold, 0.0)
-    return x1 * x2
-
-
-def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
-    d = x.shape[-1] // 2
-    return F.silu(x[..., :d]) * x[..., d:]
-
-
-def mul_and_silu(x: torch.Tensor) -> torch.Tensor:
-    d = x.shape[-1] // 2
-    return x[..., :d] * F.silu(x[..., d:])
-
-
-def gelu_and_mul(x: torch.Tensor, approximate: str) -> torch.Tensor:
-    d = x.shape[-1] // 2
-    return F.gelu(x[..., :d], approximate=approximate) * x[..., d:]
-
-def gelu(x: torch.Tensor) -> torch.Tensor:
-    return F.gelu(x)
-
-def gelu_tanh(x: torch.Tensor) -> torch.Tensor:
-    return F.gelu(x, approximate="tanh")
-
-def silu(x: torch.Tensor) -> torch.Tensor:
-    return F.silu(x)
-
-@pytest.mark.parametrize(
-    "activation_name", ["silu_and_mul", "mul_and_silu", "gelu", "gelu_tanh", "fatrelu"]
-)
-@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
-@pytest.mark.parametrize("d", D)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
-@torch.inference_mode()
-def test_act_and_mul(
-    activation_name: str,
-    num_tokens: int,
-    d: int,
-    dtype: torch.dtype,
-    seed: int,
-    device: str,
-) -> None:
-    random.seed(seed)
-    torch.manual_seed(seed)
-    torch.set_default_device(device)
-    x = torch.randn(num_tokens, 2 * d, dtype=dtype)
-    if activation_name == "silu_and_mul":
-        torch_fn = silu_and_mul
-        fn = activation.silu_and_mul
-        op = activation.ops.silu_and_mul
-        layer = activation.layers.SiluAndMul()
-    elif activation_name == "mul_and_silu":
-        torch_fn = mul_and_silu
-        fn = activation.mul_and_silu
-        op = activation.ops.mul_and_silu
-        layer = activation.layers.MulAndSilu()
-    elif activation_name == "gelu":
-        torch_fn = lambda x: gelu_and_mul(x, "none")
-        fn = activation.gelu_and_mul
-        op = activation.ops.gelu_and_mul
-        layer = activation.layers.GeluAndMul()
-    elif activation_name == "gelu_tanh":
-        torch_fn = lambda x: gelu_and_mul(x, "tanh")
-        fn = activation.gelu_tanh_and_mul
-        op = activation.ops.gelu_tanh_and_mul
-        layer = activation.layers.GeluTanhAndMul()
-    elif activation_name == "fatrelu":
-        threshold = random.uniform(0, 1)
-        torch_fn = lambda x: fatrelu_and_mul(x, threshold)
-        fn = lambda out, x: activation.fatrelu_and_mul(out, x, threshold)
-        op = activation.ops.fatrelu_and_mul
-        layer = activation.layers.FatreluAndMul(threshold)
-
-    out_shape = x.shape[:-1] + (x.shape[-1] // 2,)
-    out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
-    out = fn(out, x)
-    mod_out = layer(x)
-    ref_out = torch_fn(x)
-
-    # The SiLU, GELU and FatReLU implementations are equivalent to the native
-    # PyTorch implementations, so we can do exact comparison.
-    torch.testing.assert_close(out, ref_out, atol=0.0, rtol=0.0)
-    torch.testing.assert_close(mod_out, ref_out, atol=0.0, rtol=0.0)
-
-    d = x.shape[-1] // 2
-    output_shape = x.shape[:-1] + (d,)
-    out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-    if activation_name == "fatrelu":
-        opcheck(op, (out, x, threshold))
-    else:
-        opcheck(op, (out, x))
-
-
-@pytest.mark.parametrize(
-    "activation_fns",
-    [
-        (
-            gelu_fast,
-            activation.gelu_fast,
-            activation.ops.gelu_fast,
-            activation.layers.FastGELU,
-        ),
-        (
-            gelu_new,
-            activation.gelu_new,
-            activation.ops.gelu_new,
-            activation.layers.NewGELU,
-        ),
-        (
-            gelu_quick,
-            activation.gelu_quick,
-            activation.ops.gelu_quick,
-            activation.layers.QuickGELU,
-        ),
-        (
-            gelu_tanh,
-            activation.gelu_tanh,
-            activation.ops.gelu_tanh,
-            activation.layers.GeluTanh,
-        ),
-        (
-            silu,
-            activation.silu,
-            activation.ops.silu,
-            activation.layers.Silu,
-        ),
-        (
-            gelu, 
-            activation.gelu, 
-            activation.ops.gelu, 
-            activation.layers.Gelu
-        ),
-    ],
-)
-@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
-@pytest.mark.parametrize("d", D)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
-@torch.inference_mode()
-def test_activation(
-    activation_fns,
-    num_tokens: int,
-    d: int,
-    dtype: torch.dtype,
-    seed: int,
-    device: str,
-) -> None:
-    torch.manual_seed(seed)
-    torch.set_default_device(device)
-    x = torch.randn(num_tokens, d, dtype=dtype)
-    torch_fn, fn, op, cls = activation_fns
-    layer = cls()
-    out = fn(torch.empty_like(x), x)
-    layer_out = layer(x)
-    ref_out = torch_fn(x)
-    torch.testing.assert_close(
-        out, ref_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
-    )
-    torch.testing.assert_close(
-        out, layer_out, atol=get_default_atol(out), rtol=get_default_rtol(out)
-    )
-
-    out = torch.empty_like(x)
-    opcheck(op, (out, x))

tests/kernels/utils.py

@@ -1,73 +0,0 @@
-"""Kernel test utils"""
-
-import itertools
-import random
-import unittest
-from numbers import Number
-from typing import Any, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union
-
-import pytest
-import torch
-from torch._prims_common import TensorLikeType
-
-# For now, disable "test_aot_dispatch_dynamic" since there are some
-# bugs related to this test in PyTorch 2.4.
-DEFAULT_OPCHECK_TEST_UTILS: Tuple[str, ...] = (
-    "test_schema",
-    "test_autograd_registration",
-    "test_faketensor",
-)
-
-ALL_OPCHECK_TEST_UTILS: Tuple[str, ...] = (
-    "test_schema",
-    "test_autograd_registration",
-    "test_faketensor",
-    "test_aot_dispatch_dynamic",
-)
-
-
-# Copied/modified from torch._refs.__init__.py
-def fp8_allclose(
-    a: TensorLikeType,
-    b: TensorLikeType,
-    rtol: float = 1e-05,
-    atol: float = 1e-08,
-    equal_nan: bool = False,
-) -> bool:
-    """
-    Reference implementation of torch.allclose
-    """
-    torch._refs._check_close_args(name="torch.allclose", a=a, b=b, rtol=rtol, atol=atol)
-
-    return bool(
-        torch.all(
-            torch.isclose(
-                a.double(), b.double(), rtol=rtol, atol=atol, equal_nan=equal_nan
-            )
-        ).item()
-    )
-
-
-# A special version of op check that has a restricted default set of test_utils
-# and a patched version of allclose that supports fp8 types.
-def opcheck(
-    op: Union[
-        torch._ops.OpOverload,
-        torch._ops.OpOverloadPacket,
-        torch._library.custom_ops.CustomOpDef,
-    ],
-    args: Tuple[Any, ...],
-    kwargs: Optional[Dict[str, Any]] = None,
-    *,
-    test_utils: Union[str, Sequence[str]] = ALL_OPCHECK_TEST_UTILS,
-    raise_exception: bool = True,
-    cond: bool = True
-) -> Dict[str, str]:
-    with unittest.mock.patch("torch.allclose", new=fp8_allclose):
-        return (
-            torch.library.opcheck(
-                op, args, kwargs, test_utils=test_utils, raise_exception=raise_exception
-            )
-            if cond
-            else {}
-        )

torch-ext/activation/__init__.py

@@ -1,75 +0,0 @@
-import torch
-
-from ._ops import ops
-
-from . import layers
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out
-
-
-def mul_and_silu(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.mul_and_silu(out, x)
-    return out
-
-
-def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.gelu_and_mul(out, x)
-    return out
-
-
-def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.gelu_tanh_and_mul(out, x)
-    return out
-
-
-def fatrelu_and_mul(out: torch.Tensor, x: torch.Tensor, threshold: float = 0.0) -> None:
-    ops.fatrelu_and_mul(out, x, threshold)
-    return out
-
-
-def gelu(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.gelu(out, x)
-    return out
-
-def silu(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu(out, x)
-    return out
-
-
-def gelu_tanh(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.gelu_tanh(out, x)
-    return out
-
-
-def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.gelu_fast(out, x)
-    return out
-
-
-def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.gelu_new(out, x)
-    return out
-
-
-def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.gelu_quick(out, x)
-    return out
-
-
-__all__ = [
-    "silu_and_mul",
-    "mul_and_silu",
-    "gelu_and_mul",
-    "gelu_tanh_and_mul",
-    "fatrelu_and_mul",
-    "gelu_fast",
-    "gelu_new",
-    "gelu_quick",
-    "gelu_tanh",
-    "silu",
-    "gelu",
-    "layers",
-]

torch-ext/activation/layers.py

@@ -1,179 +0,0 @@
-import torch
-import torch.nn as nn
-
-from ._ops import ops
-
-
-class SiluAndMul(nn.Module):
-    """An activation function for SwiGLU.
-
-    The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[-1] // 2.
-
-    Shapes:
-        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
-        return: (num_tokens, d) or (batch_size, seq_len, d)
-    """
-
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor):
-        d = x.shape[-1] // 2
-        output_shape = x.shape[:-1] + (d,)
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        ops.silu_and_mul(out, x)
-        return out
-
-class Silu(nn.Module):
-    """An activation function for SiLU.
-
-    The function computes x -> silu(x).
-
-    Shapes:
-        x: (num_tokens, d) or (batch_size, seq_len, d)
-        return: (num_tokens, d) or (batch_size, seq_len, d)
-    """
-
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor):
-        out = torch.empty_like(x)
-        ops.silu(out, x)
-        return out
-
-class Gelu(nn.Module):
-    """An activation function for GELU.
-
-    The function computes x -> gelu(x).
-
-    Shapes:
-        x: (num_tokens, d) or (batch_size, seq_len, d)
-        return: (num_tokens, d) or (batch_size, seq_len, d)
-    """
-
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor):
-        out = torch.empty_like(x)
-        ops.gelu(out, x)
-        return out
-
-class GeluTanh(nn.Module):
-    """An activation function for GELU with `tanh` approximation.
-
-    The function computes x -> gelu_tanh(x).
-
-    Shapes:
-        x: (num_tokens, d) or (batch_size, seq_len, d)
-        return: (num_tokens, d) or (batch_size, seq_len, d)
-    """
-
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor):
-        out = torch.empty_like(x)
-        ops.gelu_tanh(out, x)
-        return out
-
-
-class MulAndSilu(nn.Module):
-    """An activation function for SwiGLU.
-
-    The function computes x -> x[:d] * silu(x[d:]) where d = x.shape[-1] // 2.
-
-    Shapes:
-        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
-        return: (num_tokens, d) or (batch_size, seq_len, d)
-    """
-
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = x.shape[:-1] + (d,)
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        ops.mul_and_silu(out, x)
-        return out
-
-
-class GeluAndMul(nn.Module):
-    """An activation function for GeGLU.
-
-    The function computes x -> GELU(x[:d]) * x[d:] where d = x.shape[-1] // 2.
-
-    Shapes:
-        x: (batch_size, seq_len, 2 * d) or (num_tokens, 2 * d)
-        return: (batch_size, seq_len, d) or (num_tokens, d)
-    """
-
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor):
-        d = x.shape[-1] // 2
-        output_shape = x.shape[:-1] + (d,)
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        ops.gelu_and_mul(out, x)
-        return out
-
-
-class GeluTanhAndMul(nn.Module):
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor):
-        d = x.shape[-1] // 2
-        output_shape = x.shape[:-1] + (d,)
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        ops.gelu_tanh_and_mul(out, x)
-        return out
-
-
-class FatreluAndMul(nn.Module):
-    """An activation function for FATReLU.
-
-    The function computes x -> FATReLU(x[:d]) * x[d:] where
-    d = x.shape[-1] // 2.
-    This is used in openbmb/MiniCPM-S-1B-sft.
-
-    Shapes:
-        x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
-        return: (num_tokens, d) or (batch_size, seq_len, d)
-    """
-
-    can_torch_compile: bool = True
-
-    def __init__(self, threshold: float = 0.0):
-        super().__init__()
-        self.threshold = threshold
-
-    def forward(self, x: torch.Tensor):
-        d = x.shape[-1] // 2
-        output_shape = x.shape[:-1] + (d,)
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        ops.fatrelu_and_mul(out, x, self.threshold)
-        return out
-
-
-class FastGELU(nn.Module):
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        ops.gelu_fast(out, x)
-        return out
-
-
-class NewGELU(nn.Module):
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        ops.gelu_new(out, x)
-        return out
-
-
-class QuickGELU(nn.Module):
-    can_torch_compile: bool = True
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        ops.gelu_quick(out, x)
-        return out

torch-ext/torch_binding.cpp

@@ -1,52 +0,0 @@
-#include <torch/library.h>
-
-#include "registration.h"
-#include "torch_binding.h"
-
-TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
-  // Activation ops
-  // Activation function used in SwiGLU.
-  ops.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
-  ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
-
-  ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
-  ops.impl("mul_and_silu", torch::kCUDA, &mul_and_silu);
-
-  // Activation function used in GeGLU with `none` approximation.
-  ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);
-
-  // Activation function used in GeGLU with `tanh` approximation.
-  ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_tanh_and_mul", torch::kCUDA, &gelu_tanh_and_mul);
-
-  // FATReLU implementation.
-  ops.def("fatrelu_and_mul(Tensor! out, Tensor input, float threshold) -> ()");
-  ops.impl("fatrelu_and_mul", torch::kCUDA, &fatrelu_and_mul);
-
-  // GELU implementation used in GPT-2.
-  ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_new", torch::kCUDA, &gelu_new);
-
-  // Approximate GELU implementation.
-  ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_fast", torch::kCUDA, &gelu_fast);
-
-  // Quick GELU implementation.
-  ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_quick", torch::kCUDA, &gelu_quick);
-
-  // GELU with `tanh` approximation.
-  ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu_tanh", torch::kCUDA, &gelu_tanh);
-
-  // SiLU implementation.
-  ops.def("silu(Tensor! out, Tensor input) -> ()");
-  ops.impl("silu", torch::kCUDA, &silu);
-
-  // GELU with none approximation.
-  ops.def("gelu(Tensor! out, Tensor input) -> ()");
-  ops.impl("gelu", torch::kCUDA, &gelu);
-}
-
-REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -1,26 +0,0 @@
-#pragma once
-
-#include <torch/torch.h>
-
-void silu_and_mul(torch::Tensor &out, torch::Tensor &input);
-
-void mul_and_silu(torch::Tensor& out, torch::Tensor& input);
-
-void gelu_and_mul(torch::Tensor &out, torch::Tensor &input);
-
-void gelu_tanh_and_mul(torch::Tensor &out, torch::Tensor &input);
-
-void fatrelu_and_mul(torch::Tensor &out, torch::Tensor &input,
-                     double threshold);
-
-void gelu_new(torch::Tensor &out, torch::Tensor &input);
-
-void gelu_fast(torch::Tensor &out, torch::Tensor &input);
-
-void gelu_quick(torch::Tensor &out, torch::Tensor &input);
-
-void gelu_tanh(torch::Tensor &out, torch::Tensor &input);
-
-void silu(torch::Tensor &out, torch::Tensor &input);
-
-void gelu(torch::Tensor &out, torch::Tensor &input);