Update README.md

README.md

@@ -2,9 +2,6 @@
 tags:
 - 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/activation_kernels.cu

@@ -9,16 +9,8 @@
 
 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>
+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]
@@ -27,7 +19,7 @@ __global__ void act_and_mul_kernel(
   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);
+    out[token_idx * d + idx] = ACT_FN(x) * y;
   }
 }
 
@@ -63,21 +55,16 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
 }  // 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)                 \
+#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));                                         \
-  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>  \
+        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); \
       });
@@ -85,27 +72,19 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
 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);
+  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, true);
+  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, true);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel);
 }
 
 namespace vllm {

activation/cuda_compat.h

@@ -4,10 +4,10 @@
   #include <hip/hip_runtime.h>
 #endif
 
-#if defined(USE_ROCM) && defined(__GFX9__)
-  #define WARP_SIZE 64
-#else
+#ifndef USE_ROCM
   #define WARP_SIZE 32
+#else
+  #define WARP_SIZE warpSize
 #endif
 
 #ifndef USE_ROCM

activation/dispatch_utils.h

@@ -6,11 +6,6 @@
 
 #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__)  \
@@ -19,35 +14,6 @@
 #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__)     \
@@ -65,19 +31,5 @@
   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 +1,17 @@
 [general]
 name = "activation"
-universal = false
 
 [torch]
 src = [
-    "torch-ext/torch_binding.cpp",
-    "torch-ext/torch_binding.h",
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h"
 ]
 
 [kernel.activation]
-backend = "cuda"
-depends = ["torch"]
+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",
+  "activation/activation_kernels.cu",
+  "activation/cuda_compat.h",
+  "activation/dispatch_utils.h",
 ]
+depends = [ "torch" ]

build/{torch27-cxx11-cu118-x86_64-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/{torch26-cxx11-cu118-x86_64-linux/activation/_activation_be5bedb.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:9b6ba32ecc6fc898df3b0cebee85e9afc6881749fe58142280f051ca3332d913
-size 2546864
+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_be5bedb_dirty
-ops = torch.ops._activation_be5bedb_dirty
+from . import _activation_o63kkyjirmkf4
+ops = torch.ops._activation_o63kkyjirmkf4
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_be5bedb_dirty::{op_name}"
+    return f"_activation_o63kkyjirmkf4::{op_name}"

build/{torch27-cxx11-cu126-x86_64-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/{torch26-cxx11-cu126-x86_64-linux/activation/_activation_be5bedb.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:1ce11492b9675a44afb3b896ed80e425f2a47e29481c4aad9c4a6ac59520f011
-size 2621472
+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_be5bedb_dirty
-ops = torch.ops._activation_be5bedb_dirty
+from . import _activation_vrl36m2ejer54
+ops = torch.ops._activation_vrl36m2ejer54
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_be5bedb_dirty::{op_name}"
+    return f"_activation_vrl36m2ejer54::{op_name}"

build/{torch27-cxx11-cu128-x86_64-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
@@ -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/{torch26-cxx11-cu124-x86_64-linux/activation/_activation_be5bedb.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:331dcb3900d5e47a11d3577cdbac54f15a0b6e14910239293323c1d9e4eb9f49
-size 2616928
+oid sha256:8353447f64e7d2df1a6a341d9c53bced53abef267f079923ae774170d0d57c53
+size 2427936

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

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

build/{torch28-cxx11-cu126-aarch64-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
@@ -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/{torch26-cxx98-cu118-x86_64-linux/activation/_activation_be5bedb.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:434bd1ae43b7cbdb10d86b82da9a237ec05ef9d9fb4fc15cdc9096d3d5ed3fa7
-size 2539352
+oid sha256:df184a6315118d787a1bd6b435cb45f1ca7828445a1f1c0e55c57645cfbba43a
+size 2362600

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

@@ -1,9 +1,9 @@
 import torch
-from . import _activation_0c3eb4e_dirty
-ops = torch.ops._activation_0c3eb4e_dirty
+from . import _activation_qr3gs3eckeig4
+ops = torch.ops._activation_qr3gs3eckeig4
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_activation_0c3eb4e_dirty::{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

@@ -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/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

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

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

@@ -1,128 +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 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/torch26-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
@@ -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/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

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

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

@@ -1,128 +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 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/torch26-cxx11-cu126-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/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

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

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

@@ -1,128 +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 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/torch26-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
@@ -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/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

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

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

@@ -1,128 +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 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/torch26-cxx98-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
@@ -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/torch26-cxx98-cu124-x86_64-linux/activation/_activation_be5bedb.abi3.so

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

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

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

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

@@ -1,128 +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 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/torch26-cxx98-cu126-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/torch26-cxx98-cu126-x86_64-linux/activation/_activation_be5bedb.abi3.so

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

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

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

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

@@ -1,128 +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 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