[CPU] enable fused_sigmoid_mul on CPU device (#29378)

sgl-kernel/csrc/cpu/activation.cpp

@@ -52,6 +52,55 @@ void act_and_mul_kernel_impl(
   });
 }
 
+// input : [num_tokens, dim] contiguous
+// gate : [num_tokens, num_heads, head_dim] 2d or 3d, maybe strided
+template <typename scalar_t>
+void fused_sigmoid_mul_kernel_impl(
+    scalar_t* __restrict__ output,
+    const scalar_t* __restrict__ input,
+    const scalar_t* __restrict__ gate,
+    int64_t num_tokens,
+    int64_t dim,
+    int64_t num_heads,
+    int64_t head_dim,
+    int64_t g_strideT,
+    int64_t g_strideH) {
+  using bVec = at::vec::Vectorized<scalar_t>;
+  using fVec = at::vec::Vectorized<float>;
+
+  constexpr int64_t kVecSize = bVec::size();
+  const fVec one = fVec(1.f);
+  at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
+    for (int64_t i = begin; i < end; ++i) {
+      const scalar_t* __restrict__ i_ptr = input + i * dim;
+      const scalar_t* __restrict__ g_ptr = gate + i * g_strideT;
+      scalar_t* __restrict__ o_ptr = output + i * dim;
+
+      for (int64_t h = 0; h < num_heads; ++h) {
+        const scalar_t* __restrict__ attn_ptr = i_ptr + h * head_dim;
+        const scalar_t* __restrict__ gate_ptr = g_ptr + h * g_strideH;
+        scalar_t* __restrict__ out_ptr = o_ptr + h * head_dim;
+
+        int64_t d = 0;
+#pragma GCC unroll 4
+        for (; d <= head_dim - kVecSize; d += kVecSize) {
+          auto [x_fvec0, x_fvec1] = load_float_vec2(attn_ptr + d);
+          auto [g_fvec0, g_fvec1] = load_float_vec2(gate_ptr + d);
+          x_fvec0 = x_fvec0 / (one + g_fvec0.neg().exp_u20());
+          x_fvec1 = x_fvec1 / (one + g_fvec1.neg().exp_u20());
+          convert_from_float_ext<scalar_t>(x_fvec0, x_fvec1).store(out_ptr + d);
+        }
+#pragma GCC unroll 4
+        for (; d < head_dim; ++d) {
+          float x_val = static_cast<float>(attn_ptr[d]);
+          float g_val = static_cast<float>(gate_ptr[d]);
+          out_ptr[d] = static_cast<scalar_t>(x_val / (1.f + std::exp(-g_val)));
+        }
+      }
+    }
+  });
+}
+
 }  // anonymous namespace
 
 // input   : {num_tokens, 2 * d}
@@ -130,3 +179,41 @@ at::Tensor gelu_and_mul_cpu(const at::Tensor& input) {
 
   return out;
 }
+
+at::Tensor fused_sigmoid_mul_cpu(at::Tensor& input, const at::Tensor& gate, bool inplace) {
+  CHECK_DIM(2, input);
+  const int64_t gate_dim = gate.dim();
+  TORCH_CHECK(gate_dim == 2 || gate_dim == 3, "gate must be a 2D or 3D tensor");
+  CHECK_CONTIGUOUS(input);
+  CHECK_LAST_DIM_CONTIGUOUS_INPUT(gate);
+
+  const auto st = input.scalar_type();
+  CHECK_EQ(gate.scalar_type(), st);
+
+  int64_t num_tokens = input.size(0);
+  int64_t d = input.size(1);
+
+  const bool is_gate_3d = gate_dim == 3;
+  int64_t num_heads = is_gate_3d ? gate.size(1) : 1;
+  int64_t head_dim = gate.size(-1);
+  CHECK_EQ(gate.size(0), num_tokens);
+  CHECK_EQ(d, num_heads * head_dim);
+
+  int64_t g_strideT = gate.stride(0);
+  int64_t g_strideH = is_gate_3d ? gate.stride(1) : 0;
+
+  at::Tensor out = inplace ? input : at::empty_like(input);
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_sigmoid_mul", [&] {
+    fused_sigmoid_mul_kernel_impl<scalar_t>(
+        out.data_ptr<scalar_t>(),
+        input.data_ptr<scalar_t>(),
+        gate.data_ptr<scalar_t>(),
+        num_tokens,
+        d,
+        num_heads,
+        head_dim,
+        g_strideT,
+        g_strideH);
+  });
+  return out;
+}