[WS1] Add CUDA and Triton SiLU/SwiGLU activation kernels

csrc/cuda/activation.cu

@@ -0,0 +1,253 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+//
+// Batch-invariant SiLU / SwiGLU CUDA kernels (WS1 elementwise activations).
+//
+// Semantics match NativeSiLUOp / NativeSwiGLUOp:
+//   silu(x)      = x * sigmoid(x)                 (math in fp32)
+//   swiglu(g, u) = silu(g) * u                    (math in fp32)
+//
+// Pure elementwise / token-local: no cross-row reduction, so batch size and
+// padding cannot change a row's result (Axis-A bitwise invariance).
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+
+namespace {
+
+__device__ __forceinline__ float silu_f32(float x) {
+  // sigmoid(x) = 1 / (1 + exp(-x)); use expf for device fp32.
+  const float s = 1.0f / (1.0f + expf(-x));
+  return x * s;
+}
+
+__device__ __forceinline__ float silu_grad_f32(float x) {
+  // d/dx [x * s] = s + x * s * (1 - s) = s * (1 + x * (1 - s)), s = sigmoid(x)
+  const float s = 1.0f / (1.0f + expf(-x));
+  return s * (1.0f + x * (1.0f - s));
+}
+
+template <typename scalar_t>
+__global__ void silu_forward_kernel(
+    const scalar_t* __restrict__ x,
+    scalar_t* __restrict__ y,
+    const int64_t n) {
+  const int64_t idx = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
+  if (idx >= n) {
+    return;
+  }
+  const float xv = static_cast<float>(x[idx]);
+  y[idx] = static_cast<scalar_t>(silu_f32(xv));
+}
+
+template <typename scalar_t>
+__global__ void silu_backward_kernel(
+    const scalar_t* __restrict__ dy,
+    const scalar_t* __restrict__ x,
+    scalar_t* __restrict__ dx,
+    const int64_t n) {
+  const int64_t idx = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
+  if (idx >= n) {
+    return;
+  }
+  const float dyv = static_cast<float>(dy[idx]);
+  const float xv = static_cast<float>(x[idx]);
+  dx[idx] = static_cast<scalar_t>(dyv * silu_grad_f32(xv));
+}
+
+template <typename scalar_t>
+__global__ void swiglu_forward_kernel(
+    const scalar_t* __restrict__ gate,
+    const scalar_t* __restrict__ up,
+    scalar_t* __restrict__ y,
+    const int64_t n) {
+  const int64_t idx = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
+  if (idx >= n) {
+    return;
+  }
+  const float gv = static_cast<float>(gate[idx]);
+  const float uv = static_cast<float>(up[idx]);
+  y[idx] = static_cast<scalar_t>(silu_f32(gv) * uv);
+}
+
+template <typename scalar_t>
+__global__ void swiglu_backward_kernel(
+    const scalar_t* __restrict__ dy,
+    const scalar_t* __restrict__ gate,
+    const scalar_t* __restrict__ up,
+    scalar_t* __restrict__ d_gate,
+    scalar_t* __restrict__ d_up,
+    const int64_t n) {
+  const int64_t idx = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
+  if (idx >= n) {
+    return;
+  }
+  const float dyv = static_cast<float>(dy[idx]);
+  const float gv = static_cast<float>(gate[idx]);
+  const float uv = static_cast<float>(up[idx]);
+  const float s = silu_f32(gv);
+  // d_up = dy * silu(gate); d_gate = dy * up * silu'(gate)
+  d_up[idx] = static_cast<scalar_t>(dyv * s);
+  d_gate[idx] = static_cast<scalar_t>(dyv * uv * silu_grad_f32(gv));
+}
+
+static void launch_1d(int64_t n, int& threads, int64_t& blocks) {
+  threads = 256;
+  blocks = (n + threads - 1) / threads;
+  if (blocks == 0) {
+    blocks = 1;
+  }
+}
+
+static void check_cuda_contig(const torch::Tensor& t, const char* name) {
+  TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+  // Supported activation dtypes only: fp16 / bf16 / fp32 (reject float64).
+  TORCH_CHECK(
+      t.scalar_type() == at::kHalf || t.scalar_type() == at::kBFloat16 ||
+          t.scalar_type() == at::kFloat,
+      name,
+      " must be fp16, bf16, or fp32, got ",
+      t.scalar_type());
+}
+
+static void check_same_device(
+    const torch::Tensor& lhs,
+    const torch::Tensor& rhs,
+    const char* lhs_name,
+    const char* rhs_name) {
+  TORCH_CHECK(
+      lhs.device() == rhs.device(),
+      lhs_name,
+      " and ",
+      rhs_name,
+      " must be on the same CUDA device, got ",
+      lhs.device(),
+      " and ",
+      rhs.device());
+}
+
+}  // namespace
+
+torch::Tensor silu_forward_cuda(torch::Tensor x) {
+  check_cuda_contig(x, "x");
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+  auto y = torch::empty_like(x);
+  const int64_t n = x.numel();
+  if (n == 0) {
+    return y;
+  }
+  int threads = 0;
+  int64_t blocks = 0;
+  launch_1d(n, threads, blocks);
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "silu_forward_cuda", [&] {
+        silu_forward_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
+            x.data_ptr<scalar_t>(), y.data_ptr<scalar_t>(), n);
+      });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return y;
+}
+
+torch::Tensor silu_backward_cuda(torch::Tensor dy, torch::Tensor x) {
+  check_cuda_contig(dy, "dy");
+  check_cuda_contig(x, "x");
+  check_same_device(dy, x, "dy", "x");
+  TORCH_CHECK(dy.sizes() == x.sizes(), "dy and x must share shape");
+  TORCH_CHECK(dy.scalar_type() == x.scalar_type(), "dy and x must share dtype");
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+  auto dx = torch::empty_like(x);
+  const int64_t n = x.numel();
+  if (n == 0) {
+    return dx;
+  }
+  int threads = 0;
+  int64_t blocks = 0;
+  launch_1d(n, threads, blocks);
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "silu_backward_cuda", [&] {
+        silu_backward_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
+            dy.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(), dx.data_ptr<scalar_t>(), n);
+      });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return dx;
+}
+
+torch::Tensor swiglu_forward_cuda(torch::Tensor gate, torch::Tensor up) {
+  check_cuda_contig(gate, "gate");
+  check_cuda_contig(up, "up");
+  check_same_device(gate, up, "gate", "up");
+  TORCH_CHECK(gate.sizes() == up.sizes(), "gate and up must share shape");
+  TORCH_CHECK(gate.scalar_type() == up.scalar_type(), "gate and up must share dtype");
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(gate));
+  auto y = torch::empty_like(gate);
+  const int64_t n = gate.numel();
+  if (n == 0) {
+    return y;
+  }
+  int threads = 0;
+  int64_t blocks = 0;
+  launch_1d(n, threads, blocks);
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half,
+      at::ScalarType::BFloat16,
+      gate.scalar_type(),
+      "swiglu_forward_cuda",
+      [&] {
+        swiglu_forward_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
+            gate.data_ptr<scalar_t>(), up.data_ptr<scalar_t>(), y.data_ptr<scalar_t>(), n);
+      });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return y;
+}
+
+std::vector<torch::Tensor> swiglu_backward_cuda(
+    torch::Tensor dy,
+    torch::Tensor gate,
+    torch::Tensor up) {
+  check_cuda_contig(dy, "dy");
+  check_cuda_contig(gate, "gate");
+  check_cuda_contig(up, "up");
+  check_same_device(dy, gate, "dy", "gate");
+  check_same_device(gate, up, "gate", "up");
+  TORCH_CHECK(gate.sizes() == up.sizes(), "gate and up must share shape");
+  TORCH_CHECK(dy.sizes() == gate.sizes(), "dy and gate must share shape");
+  TORCH_CHECK(dy.scalar_type() == gate.scalar_type(), "dy and gate must share dtype");
+  TORCH_CHECK(up.scalar_type() == gate.scalar_type(), "up and gate must share dtype");
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(gate));
+  auto d_gate = torch::empty_like(gate);
+  auto d_up = torch::empty_like(up);
+  const int64_t n = gate.numel();
+  if (n == 0) {
+    return {d_gate, d_up};
+  }
+  int threads = 0;
+  int64_t blocks = 0;
+  launch_1d(n, threads, blocks);
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      at::ScalarType::Half,
+      at::ScalarType::BFloat16,
+      gate.scalar_type(),
+      "swiglu_backward_cuda",
+      [&] {
+        swiglu_backward_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
+            dy.data_ptr<scalar_t>(),
+            gate.data_ptr<scalar_t>(),
+            up.data_ptr<scalar_t>(),
+            d_gate.data_ptr<scalar_t>(),
+            d_up.data_ptr<scalar_t>(),
+            n);
+      });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return {d_gate, d_up};
+}

csrc/ops.cpp

@@ -92,6 +92,15 @@ torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torc
 torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b);
 torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b);
 torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc);
+// SiLU / SwiGLU Declarations (elementwise activation, general CUDA)
+torch::Tensor silu_forward_cuda(torch::Tensor x);
+torch::Tensor silu_backward_cuda(torch::Tensor dy, torch::Tensor x);
+torch::Tensor swiglu_forward_cuda(torch::Tensor gate, torch::Tensor up);
+std::vector<torch::Tensor> swiglu_backward_cuda(
+    torch::Tensor dy,
+    torch::Tensor gate,
+    torch::Tensor up);
+
 // RMSNorm Declarations & Wrappers
 
 void rmsnorm_forward_cuda(
@@ -203,6 +212,26 @@ torch::Tensor rmsnorm_backward_dw(
   return dw;
 }
 
+// SiLU / SwiGLU wrappers (WS1 elementwise activations)
+torch::Tensor silu_forward(torch::Tensor x) {
+  return silu_forward_cuda(x);
+}
+
+torch::Tensor silu_backward(torch::Tensor dy, torch::Tensor x) {
+  return silu_backward_cuda(dy, x);
+}
+
+torch::Tensor swiglu_forward(torch::Tensor gate, torch::Tensor up) {
+  return swiglu_forward_cuda(gate, up);
+}
+
+std::vector<torch::Tensor> swiglu_backward(
+    torch::Tensor dy,
+    torch::Tensor gate,
+    torch::Tensor up) {
+  return swiglu_backward_cuda(dy, gate, up);
+}
+
 // Deterministic standard-softmax attention (issue #147)
 std::vector<torch::Tensor> deterministic_attention_forward(
     torch::Tensor q,
@@ -338,6 +367,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("rmsnorm_backward_dx", &rmsnorm_backward_dx, "Batch-invariant RMSNorm backward dx CUDA");
     m.def("rmsnorm_backward_dw", &rmsnorm_backward_dw, "Deterministic RMSNorm backward dweight CUDA");
 
+    // registry SiLU / SwiGLU (elementwise activation)
+    m.def("silu_forward", &silu_forward, "Batch-invariant SiLU forward CUDA");
+    m.def("silu_backward", &silu_backward, "Batch-invariant SiLU backward CUDA");
+    m.def("swiglu_forward", &swiglu_forward, "Batch-invariant SwiGLU forward CUDA");
+    m.def("swiglu_backward", &swiglu_backward, "Batch-invariant SwiGLU backward CUDA");
+
     // Deterministic standard-softmax attention (issue #147)
     m.def(
         "deterministic_attention_forward",