Update README.md

README.md

@@ -1,13 +1,7 @@
 ---
 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).
-
-Kernel source: https://github.com/huggingface/kernels-community/tree/main/activation
-
+Activation kernels from [vLLM](https://github.com/vllm-project/vllm/blob/main/csrc/activation_kernels.cu).

activation/activation_kernels.cu

@@ -0,0 +1,204 @@
+#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 {
+
+// Activation and gating kernel template.
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&)>
+__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] = ACT_FN(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.
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                            \
+  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", [&] {                   \
+        vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>>             \
+            <<<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);
+}
+
+void gelu_and_mul(torch::Tensor& out,    // [..., d]
+                  torch::Tensor& input)  // [..., 2 * d]
+{
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel);
+}
+
+void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
+                       torch::Tensor& input)  // [..., 2 * d]
+{
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel);
+}
+
+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);
+}

activation/cuda_compat.h

@@ -0,0 +1,49 @@
+#pragma once
+
+#ifdef USE_ROCM
+  #include <hip/hip_runtime.h>
+#endif
+
+#ifndef USE_ROCM
+  #define WARP_SIZE 32
+#else
+  #define WARP_SIZE warpSize
+#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

@@ -0,0 +1,35 @@
+/*
+ * Adapted from
+ * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h
+ */
+#pragma once
+
+#include <torch/all.h>
+
+#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__))
+
+#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_INTEGRAL_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))

build.toml

@@ -0,0 +1,17 @@
+[general]
+name = "activation"
+
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h"
+]
+
+[kernel.activation]
+cuda-capabilities = [ "7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0" ]
+src = [
+  "activation/activation_kernels.cu",
+  "activation/cuda_compat.h",
+  "activation/dispatch_utils.h",
+]
+depends = [ "torch" ]

build/{torch28-cxx11-cu126-aarch64-linux → torch25-cxx11-cu118-x86_64-linux}/activation/__init__.py

@@ -1,8 +1,15 @@
 import torch
 
-from ._ops import ops
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
 
-from . import layers
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
 
 
 def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
@@ -10,11 +17,6 @@ def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
     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
@@ -43,15 +45,3 @@ def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
 def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
     ops.gelu_quick(out, x)
     return out
-
-
-__all__ = [
-    "silu_and_mul",
-    "gelu_and_mul",
-    "gelu_tanh_and_mul",
-    "fatrelu_and_mul",
-    "gelu_fast",
-    "gelu_new",
-    "gelu_quick",
-    "layers",
-]

build/{torch27-cxx11-cu118-x86_64-linux/activation/_activation_beeaae6.abi3.so → torch25-cxx11-cu118-x86_64-linux/activation/_activation_o63kkyjirmkf4.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ce06ec284ecd4ac5423d3822a60cd9eeb686d0054b38d66567de73e1137b0567
-size 2773632
+oid sha256:d50cdabfbed1df74e921ac34ff00bca0555977b14ef8082ddae7b1f30985a494
+size 2370160

build/{torch27-cxx11-cu118-x86_64-linux → torch25-cxx11-cu118-x86_64-linux}/activation/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _activation_beeaae6
-ops = torch.ops._activation_beeaae6
+from . import _activation_o63kkyjirmkf4
+ops = torch.ops._activation_o63kkyjirmkf4
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_beeaae6::{op_name}"
+    return f"_activation_o63kkyjirmkf4::{op_name}"

build/{torch28-cxx11-cu128-aarch64-linux → torch25-cxx11-cu121-x86_64-linux}/activation/__init__.py

@@ -1,8 +1,15 @@
 import torch
 
-from ._ops import ops
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
 
-from . import layers
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
 
 
 def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
@@ -10,11 +17,6 @@ def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
     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
@@ -43,15 +45,3 @@ def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
 def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
     ops.gelu_quick(out, x)
     return out
-
-
-__all__ = [
-    "silu_and_mul",
-    "gelu_and_mul",
-    "gelu_tanh_and_mul",
-    "fatrelu_and_mul",
-    "gelu_fast",
-    "gelu_new",
-    "gelu_quick",
-    "layers",
-]

build/{torch27-cxx11-cu126-x86_64-linux/activation/_activation_beeaae6.abi3.so → torch25-cxx11-cu121-x86_64-linux/activation/_activation_vrl36m2ejer54.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a529bd105aca5081398d63329e829b6b159570424cd654d3a9f275ca9a720e82
-size 2852200
+oid sha256:2bd0709ef09c8f0c18d1dc4a36c8096c59459bece61f5f5dbea95d1e73f54d44
+size 2393264

build/{torch27-cxx11-cu126-x86_64-linux → torch25-cxx11-cu121-x86_64-linux}/activation/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _activation_beeaae6
-ops = torch.ops._activation_beeaae6
+from . import _activation_vrl36m2ejer54
+ops = torch.ops._activation_vrl36m2ejer54
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_beeaae6::{op_name}"
+    return f"_activation_vrl36m2ejer54::{op_name}"

build/{torch27-cxx11-cu128-aarch64-linux → torch25-cxx11-cu124-x86_64-linux}/activation/__init__.py

@@ -1,8 +1,15 @@
 import torch
 
-from ._ops import ops
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
 
-from . import layers
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
 
 
 def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
@@ -10,11 +17,6 @@ def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
     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
@@ -30,20 +32,6 @@ def fatrelu_and_mul(out: torch.Tensor, x: torch.Tensor, threshold: float = 0.0)
     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
@@ -57,19 +45,3 @@ def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
 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",
-]

build/{torch27-cxx11-cu128-aarch64-linux/activation/_activation_320b408.abi3.so → torch25-cxx11-cu124-x86_64-linux/activation/_activation_va3moa75vw7c2.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:34bdeb9ab72686850aef0a16b225b1b956162edb2cf46cba65c5e5b92ae267ae
-size 4207000
+oid sha256:8353447f64e7d2df1a6a341d9c53bced53abef267f079923ae774170d0d57c53
+size 2427936

build/{torch27-cxx11-cu128-aarch64-linux → torch25-cxx11-cu124-x86_64-linux}/activation/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _activation_320b408
-ops = torch.ops._activation_320b408
+from . import _activation_va3moa75vw7c2
+ops = torch.ops._activation_va3moa75vw7c2
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_320b408::{op_name}"
+    return f"_activation_va3moa75vw7c2::{op_name}"

build/{torch27-cxx11-cu128-x86_64-linux → torch25-cxx98-cu118-x86_64-linux}/activation/__init__.py

@@ -1,8 +1,15 @@
 import torch
 
-from ._ops import ops
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
 
-from . import layers
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
 
 
 def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
@@ -10,11 +17,6 @@ def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
     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
@@ -30,20 +32,6 @@ def fatrelu_and_mul(out: torch.Tensor, x: torch.Tensor, threshold: float = 0.0)
     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
@@ -57,19 +45,3 @@ def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
 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",
-]

build/{torch27-cxx11-cu128-x86_64-linux/activation/_activation_beeaae6.abi3.so → torch25-cxx98-cu118-x86_64-linux/activation/_activation_qr3gs3eckeig4.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0f2cffcb6b5b9a49f03a2df46fc2ad36765676edecb468c233e78e1f5e21e206
-size 4127872
+oid sha256:df184a6315118d787a1bd6b435cb45f1ca7828445a1f1c0e55c57645cfbba43a
+size 2362600

build/{torch27-cxx11-cu128-x86_64-linux → torch25-cxx98-cu118-x86_64-linux}/activation/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _activation_beeaae6
-ops = torch.ops._activation_beeaae6
+from . import _activation_qr3gs3eckeig4
+ops = torch.ops._activation_qr3gs3eckeig4
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_beeaae6::{op_name}"
+    return f"_activation_qr3gs3eckeig4::{op_name}"

build/torch25-cxx98-cu121-x86_64-linux/activation/__init__.py

@@ -0,0 +1,47 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
+
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(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_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

build/torch25-cxx98-cu121-x86_64-linux/activation/_activation_p7gbzt25w3zg2.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccb13cfc2e45cf483e8b9f77f1760f28b48bcf185508d51b32d45bc759c4e8bb
+size 2385440

build/torch25-cxx98-cu121-x86_64-linux/activation/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _activation_p7gbzt25w3zg2
+ops = torch.ops._activation_p7gbzt25w3zg2
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_activation_p7gbzt25w3zg2::{op_name}"

build/torch25-cxx98-cu124-x86_64-linux/activation/__init__.py

@@ -0,0 +1,47 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
+
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(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_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

build/torch25-cxx98-cu124-x86_64-linux/activation/_activation_jg7yaigtn7wco.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f8048853e8cb06e8574a9a9497800d2be438f7989d79f44dcf2e0ced38a75a9
+size 2420192

build/torch25-cxx98-cu124-x86_64-linux/activation/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _activation_jg7yaigtn7wco
+ops = torch.ops._activation_jg7yaigtn7wco
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_activation_jg7yaigtn7wco::{op_name}"

build/torch26-cxx11-cu118-x86_64-linux/activation/__init__.py

@@ -0,0 +1,47 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
+
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(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_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

build/torch26-cxx11-cu118-x86_64-linux/activation/_activation_ncisyrun7guwk.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cde5439e78ba0e1aaa1937d798b214b46d38cbab8e4384b93a22239fed1a4dd4
+size 2370264

build/torch26-cxx11-cu118-x86_64-linux/activation/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _activation_ncisyrun7guwk
+ops = torch.ops._activation_ncisyrun7guwk
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_activation_ncisyrun7guwk::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/activation/__init__.py

@@ -0,0 +1,47 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
+
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(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_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

build/torch26-cxx11-cu124-x86_64-linux/activation/_activation_ochhfvlnc3vyc.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6bd20d411c51fc8729b15cab6a60c5c9185222474aa035489e1bff299d76682
+size 2428040

build/torch26-cxx11-cu124-x86_64-linux/activation/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _activation_ochhfvlnc3vyc
+ops = torch.ops._activation_ochhfvlnc3vyc
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_activation_ochhfvlnc3vyc::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/activation/__init__.py

@@ -0,0 +1,47 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
+
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(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_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

build/torch26-cxx11-cu126-x86_64-linux/activation/_activation_u6vnqubnicksq.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41c18b20c2bf8c49d2d3088a9bc1aad4293df0b57eafc9b141a9e8e595fe551a
+size 2436672

build/torch26-cxx11-cu126-x86_64-linux/activation/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _activation_u6vnqubnicksq
+ops = torch.ops._activation_u6vnqubnicksq
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_activation_u6vnqubnicksq::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/activation/__init__.py

@@ -0,0 +1,47 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
+
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(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_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

build/torch26-cxx98-cu118-x86_64-linux/activation/_activation_2vn6ty3gfqfb6.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfbcd5da358cd5cb7982d19c8880cf4db6f08b46622a7a953f755ad59e4e1492
+size 2362752

build/torch26-cxx98-cu118-x86_64-linux/activation/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _activation_2vn6ty3gfqfb6
+ops = torch.ops._activation_2vn6ty3gfqfb6
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_activation_2vn6ty3gfqfb6::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/activation/__init__.py

@@ -0,0 +1,47 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
+
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(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_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

build/torch26-cxx98-cu124-x86_64-linux/activation/_activation_myvteedxdpqc6.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1bc928823117c800904bcd3492bf1a0c65a32f6d8a842dc039f55e29831ab49
+size 2420344

build/torch26-cxx98-cu124-x86_64-linux/activation/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _activation_myvteedxdpqc6
+ops = torch.ops._activation_myvteedxdpqc6
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_activation_myvteedxdpqc6::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/activation/__init__.py

@@ -0,0 +1,47 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _activation
+
+        ops = torch.ops._activition
+    except ImportError:
+        raise e
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(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_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

build/torch26-cxx98-cu126-x86_64-linux/activation/_activation_rbswus6emrhm2.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:474727e434a9cd4ec984a6da7124992ead4ca0fefce9581d0fd503e36c065aed
+size 2424888

build/torch26-cxx98-cu126-x86_64-linux/activation/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _activation_rbswus6emrhm2
+ops = torch.ops._activation_rbswus6emrhm2
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_activation_rbswus6emrhm2::{op_name}"

build/torch27-cxx11-cu118-x86_64-linux/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",
-]

build/torch27-cxx11-cu118-x86_64-linux/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

build/torch27-cxx11-cu126-x86_64-linux/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",
-]