diff --git a/activation/activation_cpu/cpu_features.hpp b/activation/activation_cpu/cpu_features.hpp new file mode 100644 index 00000000..aa9a02b5 --- /dev/null +++ b/activation/activation_cpu/cpu_features.hpp @@ -0,0 +1,91 @@ +#pragma once + +#ifdef _MSC_VER +#include +#else +#include +#endif +#include +#include + +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 diff --git a/activation/activation_cpu/silu_and_mul_avx512.cpp b/activation/activation_cpu/silu_and_mul_avx512.cpp new file mode 100644 index 00000000..3af36737 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +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; + using Vf = Vectorized; + 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(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 *out_data = out.data_ptr(); + + // 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(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 diff --git a/activation/activation_cpu/silu_and_mul_avx512.hpp b/activation/activation_cpu/silu_and_mul_avx512.hpp new file mode 100644 index 00000000..c57242b3 --- /dev/null +++ b/activation/activation_cpu/silu_and_mul_avx512.hpp @@ -0,0 +1,14 @@ +#pragma once + +#include + +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 diff --git a/activation/activation_cpu/silu_and_mul_cpu.cpp b/activation/activation_cpu/silu_and_mul_cpu.cpp new file mode 100644 index 00000000..01e319d5 --- /dev/null +++ b/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 + +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 diff --git a/activation/activation_cpu/silu_and_mul_cpu.hpp b/activation/activation_cpu/silu_and_mul_cpu.hpp new file mode 100644 index 00000000..3fc813ff --- /dev/null +++ b/activation/activation_cpu/silu_and_mul_cpu.hpp @@ -0,0 +1,15 @@ +#pragma once + +#include + +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 diff --git a/activation/activation_cpu/silu_and_mul_cpu_torch.cpp b/activation/activation_cpu/silu_and_mul_cpu_torch.cpp new file mode 100644 index 00000000..9a7b08dd --- /dev/null +++ b/activation/activation_cpu/silu_and_mul_cpu_torch.cpp @@ -0,0 +1,14 @@ +#include + +#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); +} diff --git a/activation/build.toml b/activation/build.toml index 76f02e12..0cd3930c 100644 --- a/activation/build.toml +++ b/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", +] diff --git a/activation/flake.nix b/activation/flake.nix index f8ee7ca9..61f5baaa 100644 --- a/activation/flake.nix +++ b/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; }; } diff --git a/activation/torch-ext/torch_binding.cpp b/activation/torch-ext/torch_binding.cpp index 07847753..39a5d2e6 100644 --- a/activation/torch-ext/torch_binding.cpp +++ b/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) -> ()");