activation: Add CPU backend for silu_and_mul

activation/activation_cpu/cpu_features.hpp

@@ -0,0 +1,91 @@
+#pragma once
+
+#ifdef _MSC_VER
+#include <intrin.h>
+#else
+#include <cpuid.h>
+#endif
+#include <fstream>
+#include <string>
+
+namespace activation_cpu {
+
+// Runtime CPU feature detection. Used to pick the AVX512 path only on hardware
+// that actually supports it; otherwise the generic ATen fallback is used.
+class CPUFeatures {
+public:
+  static bool hasAVX512BF16() {
+    static bool supported = checkAVX512BF16();
+    return supported;
+  }
+
+private:
+  static bool checkAVX512() {
+#ifdef _MSC_VER
+    int cpu_info[4];
+    __cpuid(cpu_info, 0);
+    if (cpu_info[0] < 7)
+      return false;
+    __cpuidex(cpu_info, 7, 0);
+    bool avx512f = (cpu_info[1] & (1 << 16)) != 0; // EBX bit 16
+    if (!avx512f)
+      return false;
+    __cpuid(cpu_info, 1);
+    bool osxsave = (cpu_info[2] & (1 << 27)) != 0; // ECX bit 27
+    if (!osxsave)
+      return false;
+    unsigned long long xcr0 = _xgetbv(0);
+    return ((xcr0 & 0xE6ULL) == 0xE6ULL);
+#else
+    unsigned int eax, ebx, ecx, edx;
+    if (__get_cpuid_max(0, nullptr) < 7)
+      return false;
+    __cpuid_count(7, 0, eax, ebx, ecx, edx);
+    bool avx512f = (ebx & (1 << 16)) != 0; // EBX bit 16
+    if (!avx512f)
+      return false;
+    if (__get_cpuid(1, &eax, &ebx, &ecx, &edx) == 0)
+      return false;
+    bool osxsave = (ecx & (1 << 27)) != 0; // ECX bit 27
+    if (!osxsave)
+      return false;
+    unsigned int xcr0_lo = 0, xcr0_hi = 0;
+    __asm__ volatile("xgetbv" : "=a"(xcr0_lo), "=d"(xcr0_hi) : "c"(0));
+    unsigned long long xcr0 =
+        ((unsigned long long)xcr0_hi << 32) | xcr0_lo;
+    return ((xcr0 & 0xE6ULL) == 0xE6ULL);
+#endif
+  }
+
+  static bool checkAVX512BF16() {
+    if (!checkAVX512())
+      return false;
+#ifndef _MSC_VER
+    // Most robust on Linux: read the AVX512_BF16 flag from /proc/cpuinfo.
+    std::ifstream f("/proc/cpuinfo");
+    if (f) {
+      std::string line;
+      while (std::getline(f, line)) {
+        if (line.find("avx512_bf16") != std::string::npos ||
+            line.find("avx512bf16") != std::string::npos)
+          return true;
+      }
+    }
+    // Fallback: CPUID leaf 7 subleaf 1, EAX bit 5 = AVX512_BF16.
+    unsigned int eax, ebx, ecx, edx;
+    if (__get_cpuid_max(0, nullptr) < 7)
+      return false;
+    __cpuid_count(7, 1, eax, ebx, ecx, edx);
+    return (eax & (1u << 5)) != 0;
+#else
+    int cpu_info[4];
+    __cpuid(cpu_info, 0);
+    if (cpu_info[0] < 7)
+      return false;
+    __cpuidex(cpu_info, 7, 1);
+    return (cpu_info[0] & (1 << 5)) != 0; // EAX bit 5
+#endif
+  }
+};
+
+} // namespace activation_cpu

activation/activation_cpu/silu_and_mul_avx512.cpp

@@ -0,0 +1,96 @@
+// AVX512 + AVX512-BF16 implementation of SiLU-and-Mul for BF16.
+// Compile with -mavx512f -mavx512bf16 -mavx512vl -DCPU_CAPABILITY_AVX512 so that
+// at::vec::Vectorized actually lowers to AVX512 (otherwise it falls back to the
+// scalar vec_base.h implementation).
+#include "silu_and_mul_avx512.hpp"
+
+#include <ATen/cpu/vec/vec.h>
+#include <ATen/cpu/vec/functional.h>
+#include <ATen/Parallel.h>
+#include <algorithm>
+#include <cmath>
+#include <tuple>
+
+using namespace at::vec;
+
+namespace activation_cpu {
+namespace avx512 {
+
+// Core computation for a flat element range [begin, end) of the OUTPUT.
+// Math is done entirely in fp32: bf16 is converted to float ONCE on load and
+// back ONCE on store, instead of round-tripping through bf16 on every op.
+static inline void silu_and_mul_core(const at::BFloat16 *in_data,
+                                     at::BFloat16 *out_data, int64_t begin,
+                                     int64_t end, int64_t d) {
+  using Vb = Vectorized<at::BFloat16>;
+  using Vf = Vectorized<float>;
+  constexpr int kVb = Vb::size(); // 32 bf16 lanes = 2 x 16 fp32
+  const Vf one(1.0f);
+
+  int64_t token_idx = begin / d;
+  int64_t offset = begin % d;
+  int64_t i = begin;
+
+  while (i < end) {
+    int64_t elements_to_process = std::min(d - offset, end - i);
+
+    const at::BFloat16 *a_ptr = in_data + token_idx * (2 * d) + offset;
+    const at::BFloat16 *b_ptr = a_ptr + d;
+    at::BFloat16 *out_ptr_local = out_data + i;
+
+    int64_t j = 0;
+    for (; j <= elements_to_process - kVb; j += kVb) {
+      Vb a_bf = Vb::loadu(a_ptr + j);
+      Vb b_bf = Vb::loadu(b_ptr + j);
+
+      Vf a0, a1, b0, b1;
+      std::tie(a0, a1) = convert_bfloat16_float(a_bf);
+      std::tie(b0, b1) = convert_bfloat16_float(b_bf);
+
+      // silu(a) = a / (1 + exp(-a)), in fp32.
+      Vf s0 = a0 / (one + a0.neg().exp());
+      Vf s1 = a1 / (one + a1.neg().exp());
+
+      // Match the reference F.silu(x1) * x2 exactly: it rounds silu to bf16
+      // BEFORE multiplying. Round-trip silu through bf16, then multiply.
+      std::tie(s0, s1) =
+          convert_bfloat16_float(convert_float_bfloat16(s0, s1));
+
+      convert_float_bfloat16(s0 * b0, s1 * b1).store(out_ptr_local + j);
+    }
+
+    for (; j < elements_to_process; ++j) {
+      float a = a_ptr[j];
+      float b = b_ptr[j];
+      float s = static_cast<float>(at::BFloat16(a / (1.0f + std::exp(-a))));
+      out_ptr_local[j] = at::BFloat16(s * b);
+    }
+
+    i += elements_to_process;
+    token_idx++;
+    offset = 0;
+  }
+}
+
+void silu_and_mul(torch::Tensor &out, const torch::Tensor &input) {
+  int64_t d = out.size(-1);
+  int64_t total_elements = out.numel();
+
+  const at::BFloat16 *in_data = input.data_ptr<at::BFloat16>();
+  at::BFloat16 *out_data = out.data_ptr<at::BFloat16>();
+
+  // at::parallel_for runs serially when total_elements <= grain_size, which
+  // covers the decode/tiny case without any fork/join overhead.
+  int64_t num_threads = at::get_num_threads();
+  int64_t grain_size = std::max((int64_t)8192,
+                                total_elements /
+                                    std::max<int64_t>(num_threads, 1));
+
+  at::parallel_for(0, total_elements, grain_size,
+                   [&](int64_t begin, int64_t end) {
+                     silu_and_mul_core(in_data, out_data, begin, end, d);
+                   });
+}
+
+} // namespace avx512
+} // namespace activation_cpu

activation/activation_cpu/silu_and_mul_avx512.hpp

@@ -0,0 +1,14 @@
+#pragma once
+
+#include <torch/all.h>
+
+namespace activation_cpu {
+namespace avx512 {
+
+// Optimized SiLU-and-Mul for BF16 on AVX512 (+AVX512-BF16) CPUs.
+// Requires `out` and `input` to be contiguous BF16 tensors with
+// input.size(-1) == 2 * out.size(-1).
+void silu_and_mul(torch::Tensor &out, const torch::Tensor &input);
+
+} // namespace avx512
+} // namespace activation_cpu

activation/activation_cpu/silu_and_mul_cpu.cpp

@@ -0,0 +1,31 @@
+#include "silu_and_mul_cpu.hpp"
+#include "cpu_features.hpp"
+#include "silu_and_mul_avx512.hpp"
+
+#include <ATen/ATen.h>
+
+namespace activation_cpu {
+
+void silu_and_mul(torch::Tensor &out, const torch::Tensor &input) {
+  const int64_t d = out.size(-1);
+
+  const bool can_use_avx512 =
+      input.scalar_type() == at::kBFloat16 && out.scalar_type() == at::kBFloat16 &&
+      input.is_contiguous() && out.is_contiguous() &&
+      CPUFeatures::hasAVX512BF16();
+
+  if (can_use_avx512) {
+    avx512::silu_and_mul(out, input);
+    return;
+  }
+
+  // Generic ATen fallback: works for any dtype on any CPU and matches the
+  // reference F.silu(x[..., :d]) * x[..., d:] (silu is rounded to the input
+  // dtype before the multiply, just like PyTorch eager).
+  auto input_c = input.is_contiguous() ? input : input.contiguous();
+  auto x1 = input_c.narrow(-1, 0, d);
+  auto x2 = input_c.narrow(-1, d, d);
+  out.copy_(at::silu(x1) * x2);
+}
+
+} // namespace activation_cpu

activation/activation_cpu/silu_and_mul_cpu.hpp

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <torch/all.h>
+
+namespace activation_cpu {
+
+// Dispatcher: selects the AVX512-BF16 kernel when the input is contiguous BF16
+// and the CPU supports AVX512-BF16; otherwise uses a generic ATen fallback that
+// works for any dtype on any CPU.
+//
+// Semantics: out = silu(input[..., :d]) * input[..., d:], where
+// d = out.size(-1) and input.size(-1) == 2 * d.
+void silu_and_mul(torch::Tensor &out, const torch::Tensor &input);
+
+} // namespace activation_cpu

activation/activation_cpu/silu_and_mul_cpu_torch.cpp

@@ -0,0 +1,14 @@
+#include <torch/all.h>
+
+#include "silu_and_mul_cpu.hpp"
+
+// Global entry point referenced by torch-ext/torch_binding.cpp for the CPU
+// backend. Validates shapes and forwards to the CPU dispatcher.
+void silu_and_mul(torch::Tensor &out, torch::Tensor &input) {
+  TORCH_CHECK(input.device().is_cpu(), "input must be a CPU tensor");
+  TORCH_CHECK(out.device().is_cpu(), "out must be a CPU tensor");
+  TORCH_CHECK(input.size(-1) == 2 * out.size(-1),
+              "input last dim must be twice the output last dim");
+
+  activation_cpu::silu_and_mul(out, input);
+}

activation/build.toml

@@ -5,6 +5,7 @@ license = "Apache-2.0"
 backends = [
     "cuda",
     "metal",
+    "cpu",
 ]
 
 [general.hub]
@@ -32,3 +33,35 @@ src = [
     "activation/cuda_compat.h",
     "activation/dispatch_utils.h",
 ]
+
+[kernel.activation_cpu]
+backend = "cpu"
+depends = ["torch"]
+include = ["activation_cpu"]
+src = [
+    "activation_cpu/silu_and_mul_cpu_torch.cpp",
+    "activation_cpu/silu_and_mul_cpu.cpp",
+    "activation_cpu/silu_and_mul_cpu.hpp",
+    "activation_cpu/cpu_features.hpp",
+]
+
+[kernel.activation_cpu_avx512]
+backend = "cpu"
+cxx-flags = [
+    "-O3",
+    "-fopenmp",
+    "-mf16c",
+    "-mavx512f",
+    "-mavx512bf16",
+    "-mavx512vl",
+    "-mavx512dq",
+    "-mavx512bw",
+    "-DCPU_CAPABILITY_AVX512",
+    "-DCPU_CAPABILITY=AVX512",
+]
+depends = ["torch"]
+include = ["activation_cpu"]
+src = [
+    "activation_cpu/silu_and_mul_avx512.cpp",
+    "activation_cpu/silu_and_mul_avx512.hpp",
+]

activation/flake.nix

@@ -13,5 +13,15 @@
     kernel-builder.lib.genKernelFlakeOutputs {
       inherit self;
       path = ./.;
+
+      torchVersions =
+        let
+          # For CPU builds, only x86_64-linux is currently supported.
+          cpuSupported = version: system: !(version ? "cpu") || system == "x86_64-linux";
+        in
+        allVersions:
+        builtins.map (
+          version: version // { systems = builtins.filter (cpuSupported version) version.systems; }
+        ) allVersions;
     };
 }

activation/torch-ext/torch_binding.cpp

@@ -11,6 +11,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
 #elif defined(METAL_KERNEL)
   ops.impl("silu_and_mul", torch::kMPS, &silu_and_mul);
+#elif defined(CPU_KERNEL)
+  ops.impl("silu_and_mul", torch::kCPU, &silu_and_mul);
 #endif
 
   ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");