feat: fused grouped EPolyNormTanh 커널 (fwd + bwd, expert별 계수) (#28)

* feat: fused EPolyNormTanh kernel (fwd+bwd)

* epolynorm_tanh: warp-per-row persistent kernels, one weight-grad atomic per block; tests import the package

* epolynorm_tanh: grouped (per-expert) coefficients like grouped_poly_norm

- kernel takes weight [num_total_experts, 3] + int32 offsets + expert_offset;
in-kernel binary search per row range, padding rows (>= offsets[-1]) give
zero output / zero grads, per-expert weight grads warp-reduced then
atomicAdd'ed on expert change
- two 128-bit vectors per lane in flight (ILP) in both kernels
- backward always returns four tensors (grad_scores = empty [0] without scores)
so opcheck's fake/real output counts match
- python: fused_grouped_epolynorm_tanh + grouped_epolynorm_tanh_ref; the
single-expert fused_epolynorm_tanh / epolynorm_tanh_ref wrap them
- layers.EPolyNormTanh: weight [num_experts, 3] parameter
- tests: grouped forward/backward, expert_offset + NaN padding rows,
hidden_clamp, single-expert wrapper, module vs eager, opcheck; fp32
tolerance 1e-3 (--use_fast_math)
- bench: rows spread over --experts experts (default 48 = 384/EP8)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BiM99HZnmc1dRQt4dafowG

* epolynorm_tanh: back to one vector per lane (ILP pairs measured slower); bwd opcheck without AOT output compare

bf16 N=294912xD=1152 on B300: fwd 0.45 -> 0.55 ms and bwd 0.73 -> 1.02 ms with
the two-vector loop (98 regs in bwd, selects + wasted tail lanes); SASS shows
tanhf -> MUFU.TANH (5 per element), so the 16-bit forward is MUFU-bound.
The backward's atomicAdd weight-grad makes its outputs order-dependent, so
opcheck's eager-vs-AOT comparison is skipped for that op.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BiM99HZnmc1dRQt4dafowG

* style: pre-commit (yapf/isort/clang-format) on the epolynorm_tanh files

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BiM99HZnmc1dRQt4dafowG

---------

Co-authored-by: Taehyun Kim <taehyun.kim@motiftech.io>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

activation/epolynorm_tanh.cu

@@ -0,0 +1,593 @@
+// Grouped (per-expert) EPolyNormTanh — tanh-bounded elementwise PolyNorm with
+// SiTU-style caps — forward + backward in one kernel each.
+//
+//   e      = expert of the row. Rows are grouped by expert: offsets[e] is the
+//            cumsum of tokens per local expert (int32); the coefficient row is
+//            weight[e + expert_offset].
+//   g      = clamp_max(x, C)                   (C = hidden_clamp, < 0 =
+//   disabled) u      = clamp(mul, -C, C) q_k    = tanh(g^(k+1) / s_k) k = 0, 1,
+//   2 poly   = gate_beta * tanh((w_e0 q_0 + w_e1 q_1 + w_e2 q_2) / gate_beta)
+//   up     = up_beta * tanh(u / up_beta)
+//   out    = poly * up * score                 (score: optional per-row fp32)
+//
+// weight [num_total_experts, 3] fp32 holds the coefficients AFTER the sigmoid;
+// their gradient comes back as fp32 [num_total_experts, 3] and the caller's
+// autograd carries it through the sigmoid. Everything is computed in fp32 and
+// stored in the input dtype. Rows at or beyond offsets[-1] are padding: zero
+// output, zero input gradients, no contribution to the weight gradient.
+//
+// Launch shape: one WARP per row, each warp owning a CONTIGUOUS chunk of rows
+// (persistent grid). Rows of one expert are contiguous, so a warp's weight-grad
+// partials stay in registers until the expert changes (or the chunk ends) and
+// reach global memory with a few atomicAdds per warp. Per-row reductions
+// (grad_scores) are warp shuffles. Loads and stores are 128-bit vectorized when
+// D % width == 0 (width 8 for 16-bit types, 4 for fp32), scalar fallback
+// otherwise.
+//
+// Cost model (B300, bf16, N=294912 x D=1152): tanhf compiles to MUFU.TANH
+// under --use_fast_math, 5 per element in each kernel; at 16 MUFU ops/clk/SM
+// that is a ~0.36 ms floor against a ~0.27 ms HBM floor for the forward, so
+// the 16-bit forward is MUFU-bound (fp32, with twice the bytes, is
+// memory-bound). Two-vectors-per-lane ILP was measured slower (extra selects
+// and wasted lanes on the tail chunk).
+
+#include <ATen/Functions.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include <algorithm>
+
+#include "assert_utils.h"
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+
+namespace motif {
+namespace epoly {
+
+template <typename type, int N> struct alignas(sizeof(type) * N) vec_t {
+  type data[N];
+};
+
+constexpr int BLOCK = 256;
+constexpr int WARPS_PER_BLOCK = BLOCK / WARP_SIZE;
+
+struct Params {
+  float gate_beta, inv_gate_beta;
+  float up_beta, inv_up_beta;
+  float inv_s0, inv_s1, inv_s2;
+  float hidden_clamp; // < 0 = disabled
+};
+
+__device__ __forceinline__ float warp_sum(float v) {
+#ifndef USE_ROCM
+  constexpr unsigned int FULL_MASK = 0xffffffff;
+#else
+  constexpr unsigned long long FULL_MASK = 0xffffffffffffffffULL;
+#endif
+#pragma unroll
+  for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1)
+    v += __shfl_xor_sync(FULL_MASK, v, mask);
+  return v;
+}
+
+// Expert of `row`: first e with offsets[e] > row (offsets = cumsum of counts).
+__device__ __forceinline__ int find_expert(const int32_t *__restrict__ offsets,
+                                           int num_experts, int64_t row) {
+  int lo = 0, hi = num_experts;
+  while (lo < hi) {
+    const int mid = (lo + hi) >> 1;
+    if ((int64_t)offsets[mid] <= row)
+      lo = mid + 1;
+    else
+      hi = mid;
+  }
+  return lo;
+}
+
+// Forward math for one element: out = poly * up (score applied by the caller).
+__device__ __forceinline__ float fwd_elem(float x, float m, const Params &p,
+                                          float w0, float w1, float w2,
+                                          bool do_clamp) {
+  const float g = do_clamp ? fminf(x, p.hidden_clamp) : x;
+  const float u =
+      do_clamp ? fminf(fmaxf(m, -p.hidden_clamp), p.hidden_clamp) : m;
+  const float g2 = g * g;
+  const float q0 = tanhf(g * p.inv_s0);
+  const float q1 = tanhf(g2 * p.inv_s1);
+  const float q2 = tanhf(g2 * g * p.inv_s2);
+  const float poly =
+      p.gate_beta * tanhf((w0 * q0 + w1 * q1 + w2 * q2) * p.inv_gate_beta);
+  const float up = p.up_beta * tanhf(u * p.inv_up_beta);
+  return poly * up;
+}
+
+// Backward math for one element. Writes dx, dm; accumulates dw_k and the
+// grad_scores partial (go * poly * up).
+__device__ __forceinline__ void bwd_elem(float x, float m, float go,
+                                         float score, const Params &p, float w0,
+                                         float w1, float w2, bool do_clamp,
+                                         float &dx, float &dm, float &dw0,
+                                         float &dw1, float &dw2, float &gs) {
+  const float g = do_clamp ? fminf(x, p.hidden_clamp) : x;
+  const float u =
+      do_clamp ? fminf(fmaxf(m, -p.hidden_clamp), p.hidden_clamp) : m;
+  const float g2 = g * g;
+  const float q0 = tanhf(g * p.inv_s0);
+  const float q1 = tanhf(g2 * p.inv_s1);
+  const float q2 = tanhf(g2 * g * p.inv_s2);
+  const float tz = tanhf((w0 * q0 + w1 * q1 + w2 * q2) * p.inv_gate_beta);
+  const float poly = p.gate_beta * tz;
+  const float tu = tanhf(u * p.inv_up_beta);
+  const float up = p.up_beta * tu;
+
+  // d out / d poly = up * score ; d out / d up = poly * score
+  const float d_poly = go * up * score;
+  const float d_up = go * poly * score;
+  // poly = gb * tanh(S / gb) -> d poly / d S = 1 - tz^2
+  const float dS = d_poly * (1.f - tz * tz);
+  // up = ub * tanh(u / ub)   -> d up / d u = 1 - tu^2
+  float du = d_up * (1.f - tu * tu);
+  float dg = dS * (w0 * (1.f - q0 * q0) * p.inv_s0 +
+                   w1 * (1.f - q1 * q1) * (2.f * g * p.inv_s1) +
+                   w2 * (1.f - q2 * q2) * (3.f * g2 * p.inv_s2));
+  if (do_clamp) {
+    if (x > p.hidden_clamp)
+      dg = 0.f; // clamp_max(x, C): zero slope above C
+    if (fabsf(m) > p.hidden_clamp)
+      du = 0.f;
+  }
+  dx = dg;
+  dm = du;
+  dw0 += dS * q0;
+  dw1 += dS * q1;
+  dw2 += dS * q2;
+  gs += go * poly * up;
+}
+
+// Warp-reduce the weight-grad partials and add them to expert `e`'s row.
+__device__ __forceinline__ void flush_weight_grad(float &dw0, float &dw1,
+                                                  float &dw2, int e,
+                                                  float *__restrict__ wg) {
+  const float a = warp_sum(dw0), b = warp_sum(dw1), c = warp_sum(dw2);
+  if (threadIdx.x % WARP_SIZE == 0) {
+    atomicAdd(&wg[(int64_t)e * 3 + 0], a);
+    atomicAdd(&wg[(int64_t)e * 3 + 1], b);
+    atomicAdd(&wg[(int64_t)e * 3 + 2], c);
+  }
+  dw0 = dw1 = dw2 = 0.f;
+}
+
+// The contiguous row range [begin, end) of this warp.
+__device__ __forceinline__ void warp_rows(int64_t N, int64_t &begin,
+                                          int64_t &end) {
+  const int64_t total_warps = (int64_t)gridDim.x * WARPS_PER_BLOCK;
+  const int64_t gw =
+      (int64_t)blockIdx.x * WARPS_PER_BLOCK + threadIdx.x / WARP_SIZE;
+  const int64_t per_warp = (N + total_warps - 1) / total_warps;
+  begin = gw * per_warp;
+  end = begin + per_warp < N ? begin + per_warp : N;
+}
+
+// Load the expert's three coefficients.
+__device__ __forceinline__ void load_w(const float *__restrict__ weight, int e,
+                                       int expert_offset, float &w0, float &w1,
+                                       float &w2) {
+  const int64_t wi = (int64_t)(e + expert_offset) * 3;
+  w0 = weight[wi];
+  w1 = weight[wi + 1];
+  w2 = weight[wi + 2];
+}
+
+// ---------------------------------------------------------------------------
+// Forward — vectorized (width > 0)
+// ---------------------------------------------------------------------------
+template <typename scalar_t, int width>
+__global__ void __launch_bounds__(BLOCK) epolynorm_tanh_fwd_kernel(
+    scalar_t *__restrict__ output, const scalar_t *__restrict__ input,
+    const scalar_t *__restrict__ mul,
+    const float *__restrict__ weight,    // [num_total_experts, 3] fp32
+    const int32_t *__restrict__ offsets, // [num_experts]
+    const float *__restrict__ scores,    // nullable, [N] fp32
+    const Params p, const int D, const int64_t N, const int num_experts,
+    const int expert_offset) {
+  using v_t = vec_t<scalar_t, width>;
+  const bool do_clamp = (p.hidden_clamp >= 0.f);
+  const int vec_d = D / width;
+  const int lane = threadIdx.x % WARP_SIZE;
+  const int64_t n_valid = offsets[num_experts - 1];
+  int64_t row, end;
+  warp_rows(N, row, end);
+
+  v_t zero_v;
+#pragma unroll
+  for (int j = 0; j < width; ++j)
+    zero_v.data[j] = (scalar_t)0;
+
+  int e = -1;
+  float w0 = 0.f, w1 = 0.f, w2 = 0.f;
+  for (; row < end; ++row) {
+    const int64_t base = row * vec_d;
+    v_t *__restrict__ out_v = reinterpret_cast<v_t *>(output) + base;
+    if (row >= n_valid) { // padding row
+      for (int i = lane; i < vec_d; i += WARP_SIZE)
+        out_v[i] = zero_v;
+      continue;
+    }
+    if (e < 0 || row >= (int64_t)offsets[e]) {
+      e = find_expert(offsets, num_experts, row);
+      load_w(weight, e, expert_offset, w0, w1, w2);
+    }
+    const float score = (scores != nullptr) ? scores[row] : 1.f;
+    const v_t *__restrict__ in_v = reinterpret_cast<const v_t *>(input) + base;
+    const v_t *__restrict__ m_v = reinterpret_cast<const v_t *>(mul) + base;
+    for (int i = lane; i < vec_d; i += WARP_SIZE) {
+      const v_t xv = in_v[i];
+      const v_t mv = m_v[i];
+      v_t ov;
+#pragma unroll
+      for (int j = 0; j < width; ++j) {
+        const float o = fwd_elem((float)xv.data[j], (float)mv.data[j], p, w0,
+                                 w1, w2, do_clamp) *
+                        score;
+        ov.data[j] = (scalar_t)o;
+      }
+      out_v[i] = ov;
+    }
+  }
+}
+
+// Scalar fallback forward (D not a multiple of the vector width)
+template <typename scalar_t>
+__global__ void __launch_bounds__(BLOCK)
+    epolynorm_tanh_fwd_scalar(scalar_t *__restrict__ output,
+                              const scalar_t *__restrict__ input,
+                              const scalar_t *__restrict__ mul,
+                              const float *__restrict__ weight,
+                              const int32_t *__restrict__ offsets,
+                              const float *__restrict__ scores, const Params p,
+                              const int D, const int64_t N,
+                              const int num_experts, const int expert_offset) {
+  const bool do_clamp = (p.hidden_clamp >= 0.f);
+  const int lane = threadIdx.x % WARP_SIZE;
+  const int64_t n_valid = offsets[num_experts - 1];
+  int64_t row, end;
+  warp_rows(N, row, end);
+
+  int e = -1;
+  float w0 = 0.f, w1 = 0.f, w2 = 0.f;
+  for (; row < end; ++row) {
+    const int64_t off = row * D;
+    if (row >= n_valid) {
+      for (int i = lane; i < D; i += WARP_SIZE)
+        output[off + i] = (scalar_t)0;
+      continue;
+    }
+    if (e < 0 || row >= (int64_t)offsets[e]) {
+      e = find_expert(offsets, num_experts, row);
+      load_w(weight, e, expert_offset, w0, w1, w2);
+    }
+    const float score = (scores != nullptr) ? scores[row] : 1.f;
+    for (int i = lane; i < D; i += WARP_SIZE) {
+      const float o = fwd_elem((float)input[off + i], (float)mul[off + i], p,
+                               w0, w1, w2, do_clamp) *
+                      score;
+      output[off + i] = (scalar_t)o;
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Backward — vectorized (width > 0)
+// ---------------------------------------------------------------------------
+template <typename scalar_t, int width>
+__global__ void __launch_bounds__(BLOCK) epolynorm_tanh_bwd_kernel(
+    scalar_t *__restrict__ grad_input, scalar_t *__restrict__ grad_mul,
+    float *__restrict__ weight_grad, // [num_total_experts, 3] fp32, atomicAdd
+    float *__restrict__ grad_scores, // nullable, [N] fp32
+    const scalar_t *__restrict__ grad_output,
+    const scalar_t *__restrict__ input, const scalar_t *__restrict__ mul,
+    const float *__restrict__ weight, const int32_t *__restrict__ offsets,
+    const float *__restrict__ scores, const Params p, const int D,
+    const int64_t N, const int num_experts, const int expert_offset) {
+  using v_t = vec_t<scalar_t, width>;
+  const bool do_clamp = (p.hidden_clamp >= 0.f);
+  const int vec_d = D / width;
+  const int lane = threadIdx.x % WARP_SIZE;
+  const int64_t n_valid = offsets[num_experts - 1];
+  int64_t row, end;
+  warp_rows(N, row, end);
+
+  v_t zero_v;
+#pragma unroll
+  for (int j = 0; j < width; ++j)
+    zero_v.data[j] = (scalar_t)0;
+
+  int e = -1;
+  float w0 = 0.f, w1 = 0.f, w2 = 0.f;
+  float dw0 = 0.f, dw1 = 0.f, dw2 = 0.f;
+  for (; row < end; ++row) {
+    const int64_t base = row * vec_d;
+    v_t *__restrict__ gi_v = reinterpret_cast<v_t *>(grad_input) + base;
+    v_t *__restrict__ gm_v = reinterpret_cast<v_t *>(grad_mul) + base;
+    if (row >= n_valid) { // padding row
+      for (int i = lane; i < vec_d; i += WARP_SIZE) {
+        gi_v[i] = zero_v;
+        gm_v[i] = zero_v;
+      }
+      if (grad_scores != nullptr && lane == 0)
+        grad_scores[row] = 0.f;
+      continue;
+    }
+    if (e < 0 || row >= (int64_t)offsets[e]) {
+      if (e >= 0)
+        flush_weight_grad(dw0, dw1, dw2, e + expert_offset, weight_grad);
+      e = find_expert(offsets, num_experts, row);
+      load_w(weight, e, expert_offset, w0, w1, w2);
+    }
+    const float score = (scores != nullptr) ? scores[row] : 1.f;
+    const v_t *__restrict__ in_v = reinterpret_cast<const v_t *>(input) + base;
+    const v_t *__restrict__ m_v = reinterpret_cast<const v_t *>(mul) + base;
+    const v_t *__restrict__ go_v =
+        reinterpret_cast<const v_t *>(grad_output) + base;
+    float gs = 0.f;
+    for (int i = lane; i < vec_d; i += WARP_SIZE) {
+      const v_t xv = in_v[i];
+      const v_t mv = m_v[i];
+      const v_t gv = go_v[i];
+      v_t gi, gm;
+#pragma unroll
+      for (int j = 0; j < width; ++j) {
+        float dx, dm;
+        bwd_elem((float)xv.data[j], (float)mv.data[j], (float)gv.data[j], score,
+                 p, w0, w1, w2, do_clamp, dx, dm, dw0, dw1, dw2, gs);
+        gi.data[j] = (scalar_t)dx;
+        gm.data[j] = (scalar_t)dm;
+      }
+      gi_v[i] = gi;
+      gm_v[i] = gm;
+    }
+    if (grad_scores != nullptr) {
+      gs = warp_sum(gs);
+      if (lane == 0)
+        grad_scores[row] = gs;
+    }
+  }
+  if (e >= 0)
+    flush_weight_grad(dw0, dw1, dw2, e + expert_offset, weight_grad);
+}
+
+// Scalar fallback backward
+template <typename scalar_t>
+__global__ void __launch_bounds__(BLOCK) epolynorm_tanh_bwd_scalar(
+    scalar_t *__restrict__ grad_input, scalar_t *__restrict__ grad_mul,
+    float *__restrict__ weight_grad, float *__restrict__ grad_scores,
+    const scalar_t *__restrict__ grad_output,
+    const scalar_t *__restrict__ input, const scalar_t *__restrict__ mul,
+    const float *__restrict__ weight, const int32_t *__restrict__ offsets,
+    const float *__restrict__ scores, const Params p, const int D,
+    const int64_t N, const int num_experts, const int expert_offset) {
+  const bool do_clamp = (p.hidden_clamp >= 0.f);
+  const int lane = threadIdx.x % WARP_SIZE;
+  const int64_t n_valid = offsets[num_experts - 1];
+  int64_t row, end;
+  warp_rows(N, row, end);
+
+  int e = -1;
+  float w0 = 0.f, w1 = 0.f, w2 = 0.f;
+  float dw0 = 0.f, dw1 = 0.f, dw2 = 0.f;
+  for (; row < end; ++row) {
+    const int64_t off = row * D;
+    if (row >= n_valid) {
+      for (int i = lane; i < D; i += WARP_SIZE) {
+        grad_input[off + i] = (scalar_t)0;
+        grad_mul[off + i] = (scalar_t)0;
+      }
+      if (grad_scores != nullptr && lane == 0)
+        grad_scores[row] = 0.f;
+      continue;
+    }
+    if (e < 0 || row >= (int64_t)offsets[e]) {
+      if (e >= 0)
+        flush_weight_grad(dw0, dw1, dw2, e + expert_offset, weight_grad);
+      e = find_expert(offsets, num_experts, row);
+      load_w(weight, e, expert_offset, w0, w1, w2);
+    }
+    const float score = (scores != nullptr) ? scores[row] : 1.f;
+    float gs = 0.f;
+    for (int i = lane; i < D; i += WARP_SIZE) {
+      float dx, dm;
+      bwd_elem((float)input[off + i], (float)mul[off + i],
+               (float)grad_output[off + i], score, p, w0, w1, w2, do_clamp, dx,
+               dm, dw0, dw1, dw2, gs);
+      grad_input[off + i] = (scalar_t)dx;
+      grad_mul[off + i] = (scalar_t)dm;
+    }
+    if (grad_scores != nullptr) {
+      gs = warp_sum(gs);
+      if (lane == 0)
+        grad_scores[row] = gs;
+    }
+  }
+  if (e >= 0)
+    flush_weight_grad(dw0, dw1, dw2, e + expert_offset, weight_grad);
+}
+
+static int grid_for(int64_t N) {
+  // Enough warps to give each one a short contiguous row chunk, capped so the
+  // grid stays resident (16 blocks of 8 warps per SM).
+  const int sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  const int64_t blocks = (N + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK;
+  return (int)std::min<int64_t>(blocks, (int64_t)sms * 16);
+}
+
+static Params make_params(double gate_beta, double up_beta, double s0,
+                          double s1, double s2, double hidden_clamp) {
+  Params p;
+  p.gate_beta = (float)gate_beta;
+  p.inv_gate_beta = 1.f / (float)gate_beta;
+  p.up_beta = (float)up_beta;
+  p.inv_up_beta = 1.f / (float)up_beta;
+  p.inv_s0 = 1.f / (float)s0;
+  p.inv_s1 = 1.f / (float)s1;
+  p.inv_s2 = 1.f / (float)s2;
+  p.hidden_clamp = (float)hidden_clamp;
+  return p;
+}
+
+} // namespace epoly
+} // namespace motif
+
+// ---------------------------------------------------------------------------
+// Host entry points
+// ---------------------------------------------------------------------------
+static void check_inputs(const torch::Tensor &input, const torch::Tensor &mul,
+                         const torch::Tensor &weight,
+                         const torch::Tensor &offsets, int64_t expert_offset,
+                         const std::optional<torch::Tensor> &scores) {
+  TORCH_CHECK(input.is_cuda() && mul.is_cuda(), "inputs must be CUDA tensors");
+  TORCH_CHECK(input.dim() == 2 && mul.dim() == 2, "input/mul must be [N, D]");
+  TORCH_CHECK(input.sizes() == mul.sizes(), "input and mul must match");
+  TORCH_CHECK(input.scalar_type() == mul.scalar_type(),
+              "input and mul must share a dtype");
+  TORCH_CHECK(input.is_contiguous() && mul.is_contiguous(),
+              "input/mul must be contiguous");
+  TORCH_CHECK(weight.is_cuda() && weight.scalar_type() == torch::kFloat &&
+                  weight.dim() == 2 && weight.size(1) == 3 &&
+                  weight.is_contiguous(),
+              "weight must be a contiguous fp32 CUDA tensor [num_experts, 3]");
+  TORCH_CHECK(offsets.is_cuda() && offsets.scalar_type() == torch::kInt &&
+                  offsets.dim() == 1 && offsets.numel() >= 1 &&
+                  offsets.is_contiguous(),
+              "offsets must be a contiguous int32 CUDA tensor [num_experts]");
+  TORCH_CHECK(expert_offset >= 0 &&
+                  expert_offset + offsets.numel() <= weight.size(0),
+              "expert_offset + len(offsets) must fit weight's rows");
+  if (scores.has_value()) {
+    TORCH_CHECK(scores->is_cuda() && scores->scalar_type() == torch::kFloat,
+                "scores must be a fp32 CUDA tensor");
+    TORCH_CHECK(scores->numel() == input.size(0),
+                "scores must have one entry per row");
+    TORCH_CHECK(scores->is_contiguous(), "scores must be contiguous");
+  }
+}
+
+torch::Tensor
+epolynorm_tanh_forward(const torch::Tensor &input, const torch::Tensor &mul,
+                       const torch::Tensor &weight,
+                       const torch::Tensor &offsets, int64_t expert_offset,
+                       const std::optional<torch::Tensor> &scores,
+                       double gate_beta, double up_beta, double s0, double s1,
+                       double s2, double hidden_clamp) {
+  check_inputs(input, mul, weight, offsets, expert_offset, scores);
+  const int D = input.size(-1);
+  const int64_t N = input.size(0);
+  const int num_experts = offsets.numel();
+  auto output = torch::empty_like(input);
+  if (N == 0)
+    return output;
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const float *scores_ptr =
+      scores.has_value() ? scores->data_ptr<float>() : nullptr;
+  const motif::epoly::Params p =
+      motif::epoly::make_params(gate_beta, up_beta, s0, s1, s2, hidden_clamp);
+  dim3 grid(motif::epoly::grid_for(N));
+  dim3 block(motif::epoly::BLOCK);
+
+#define FWD_LAUNCH(width_val)                                                  \
+  MOTIF_DISPATCH_FLOATING_TYPES(                                               \
+      input.scalar_type(), "epolynorm_tanh_fwd", [&] {                         \
+        motif::epoly::epolynorm_tanh_fwd_kernel<scalar_t, width_val>           \
+            <<<grid, block, 0, stream>>>(                                      \
+                output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),       \
+                mul.data_ptr<scalar_t>(), weight.data_ptr<float>(),            \
+                offsets.data_ptr<int32_t>(), scores_ptr, p, D, N, num_experts, \
+                (int)expert_offset);                                           \
+      })
+  if (D % 8 == 0 && input.element_size() == 2)
+    FWD_LAUNCH(8);
+  else if (D % 4 == 0 && input.element_size() == 4)
+    FWD_LAUNCH(4);
+  else {
+    MOTIF_DISPATCH_FLOATING_TYPES(
+        input.scalar_type(), "epolynorm_tanh_fwd_scalar", [&] {
+          motif::epoly::epolynorm_tanh_fwd_scalar<scalar_t>
+              <<<grid, block, 0, stream>>>(
+                  output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
+                  mul.data_ptr<scalar_t>(), weight.data_ptr<float>(),
+                  offsets.data_ptr<int32_t>(), scores_ptr, p, D, N, num_experts,
+                  (int)expert_offset);
+        });
+  }
+#undef FWD_LAUNCH
+  return output;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+epolynorm_tanh_backward(const torch::Tensor &grad_output,
+                        const torch::Tensor &input, const torch::Tensor &mul,
+                        const torch::Tensor &weight,
+                        const torch::Tensor &offsets, int64_t expert_offset,
+                        const std::optional<torch::Tensor> &scores,
+                        double gate_beta, double up_beta, double s0, double s1,
+                        double s2, double hidden_clamp) {
+  check_inputs(input, mul, weight, offsets, expert_offset, scores);
+  TORCH_CHECK(grad_output.sizes() == input.sizes() &&
+                  grad_output.scalar_type() == input.scalar_type() &&
+                  grad_output.is_contiguous(),
+              "grad_output must match input");
+  const int D = input.size(-1);
+  const int64_t N = input.size(0);
+  const int num_experts = offsets.numel();
+  auto grad_input = torch::empty_like(input);
+  auto grad_mul = torch::empty_like(mul);
+  auto weight_grad = torch::zeros_like(weight);
+  // grad_scores: [N, 1] fp32 when scores are given, else an empty [0] tensor
+  // (never undefined, so the op always returns four tensors).
+  auto grad_scores = scores.has_value() ? torch::empty({N, 1}, weight.options())
+                                        : torch::empty({0}, weight.options());
+  if (N == 0)
+    return {grad_input, grad_mul, weight_grad, grad_scores};
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const float *scores_ptr =
+      scores.has_value() ? scores->data_ptr<float>() : nullptr;
+  float *gs_ptr = scores.has_value() ? grad_scores.data_ptr<float>() : nullptr;
+  const motif::epoly::Params p =
+      motif::epoly::make_params(gate_beta, up_beta, s0, s1, s2, hidden_clamp);
+  dim3 grid(motif::epoly::grid_for(N));
+  dim3 block(motif::epoly::BLOCK);
+
+#define BWD_LAUNCH(width_val)                                                  \
+  MOTIF_DISPATCH_FLOATING_TYPES(                                               \
+      input.scalar_type(), "epolynorm_tanh_bwd", [&] {                         \
+        motif::epoly::epolynorm_tanh_bwd_kernel<scalar_t, width_val>           \
+            <<<grid, block, 0, stream>>>(                                      \
+                grad_input.data_ptr<scalar_t>(),                               \
+                grad_mul.data_ptr<scalar_t>(), weight_grad.data_ptr<float>(),  \
+                gs_ptr, grad_output.data_ptr<scalar_t>(),                      \
+                input.data_ptr<scalar_t>(), mul.data_ptr<scalar_t>(),          \
+                weight.data_ptr<float>(), offsets.data_ptr<int32_t>(),         \
+                scores_ptr, p, D, N, num_experts, (int)expert_offset);         \
+      })
+  if (D % 8 == 0 && input.element_size() == 2)
+    BWD_LAUNCH(8);
+  else if (D % 4 == 0 && input.element_size() == 4)
+    BWD_LAUNCH(4);
+  else {
+    MOTIF_DISPATCH_FLOATING_TYPES(
+        input.scalar_type(), "epolynorm_tanh_bwd_scalar", [&] {
+          motif::epoly::epolynorm_tanh_bwd_scalar<scalar_t>
+              <<<grid, block, 0, stream>>>(
+                  grad_input.data_ptr<scalar_t>(),
+                  grad_mul.data_ptr<scalar_t>(), weight_grad.data_ptr<float>(),
+                  gs_ptr, grad_output.data_ptr<scalar_t>(),
+                  input.data_ptr<scalar_t>(), mul.data_ptr<scalar_t>(),
+                  weight.data_ptr<float>(), offsets.data_ptr<int32_t>(),
+                  scores_ptr, p, D, N, num_experts, (int)expert_offset);
+        });
+  }
+#undef BWD_LAUNCH
+  return {grad_input, grad_mul, weight_grad, grad_scores};
+}

benchmarks/bench_epolynorm_tanh.py

@@ -0,0 +1,234 @@
+"""Bandwidth benchmark for the fused grouped EPolyNormTanh kernels.
+
+Compares, per (N, D, dtype, scores):
+  - fused CUDA forward / backward (torch.ops._activation.epolynorm_tanh_*)
+  - the eager grouped reference (what llm-training runs today, per-expert
+    coefficients gathered per row) forward / backward
+  - torch.compile'd reference forward / backward
+  - fused through autograd (what training uses)
+and reports achieved GB/s against the GPU's theoretical HBM bandwidth (NVML bus
+width x memory clock, DDR) and against a measured device-copy ceiling.
+
+Rows are spread over --experts experts (random multinomial counts, like a
+routed MoE microbatch); the weight has --experts + 4 rows and expert_offset 4.
+
+Byte accounting (the minimum traffic of each op, in the input dtype):
+  forward : read x, mul; write out                         -> 3 * N * D * esz
+  backward: read x, mul, grad_out; write grad_x, grad_mul  -> 5 * N * D * esz
+(+ N * 4 for the fp32 scores on each side when scores are used)
+
+Usage: python benchmarks/bench_epolynorm_tanh.py [--d 1152] [--n 294912 ...]
+"""
+
+import argparse
+import statistics
+
+import torch
+
+torch._dynamo.config.cache_size_limit = 64
+
+from activation.epolynorm_tanh import (DEFAULT_GATE_BETA, DEFAULT_SCALES,
+                                       DEFAULT_UP_BETA,
+                                       fused_grouped_epolynorm_tanh,
+                                       grouped_epolynorm_tanh_ref)
+
+_ops = torch.ops._activation
+EXPERT_OFFSET = 4
+
+
+def theoretical_bw_gbs() -> tuple[float, str]:
+    try:
+        import pynvml as n
+        n.nvmlInit()
+        h = n.nvmlDeviceGetHandleByIndex(torch.cuda.current_device())
+        bus = n.nvmlDeviceGetMemoryBusWidth(h)
+        clk = n.nvmlDeviceGetMaxClockInfo(h, n.NVML_CLOCK_MEM)  # MHz
+        name = n.nvmlDeviceGetName(h)
+        # HBM3/3e: NVML reports the memory clock; data rate is 2x (DDR).
+        gbs = 2.0 * clk * 1e6 * bus / 8 / 1e9
+        return gbs, f"{name}: bus {bus} bit x {clk} MHz x2 (NVML)"
+    except Exception as e:  # pragma: no cover
+        name = torch.cuda.get_device_name()
+        table = {
+            "B300": 8000.0,
+            "B200": 8000.0,
+            "H100": 3350.0,
+            "H200": 4800.0
+        }
+        for k, v in table.items():
+            if k in name:
+                return v, f"{name}: spec table ({e.__class__.__name__})"
+        return float("nan"), f"{name}: unknown ({e})"
+
+
+def bench(fn, iters=50, warmup=10) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    times = []
+    for _ in range(iters):
+        s = torch.cuda.Event(enable_timing=True)
+        e = torch.cuda.Event(enable_timing=True)
+        s.record()
+        fn()
+        e.record()
+        torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    return statistics.median(times)
+
+
+def copy_ceiling_gbs(nbytes: int) -> float:
+    x = torch.empty(nbytes // 2, dtype=torch.bfloat16, device="cuda")
+    y = torch.empty_like(x)
+    ms = bench(lambda: y.copy_(x))
+    return 2 * x.nbytes / ms / 1e6
+
+
+def make_offsets(n, num_experts):
+    torch.manual_seed(0)
+    assign = torch.multinomial(torch.ones(num_experts) / num_experts,
+                               n,
+                               replacement=True)
+    counts = torch.bincount(assign, minlength=num_experts)
+    return torch.cumsum(counts, 0).to(torch.int32).cuda()
+
+
+def run(n, d, num_experts, dtype, with_scores, peak):
+    esz = torch.tensor([], dtype=dtype).element_size()
+    x = torch.randn(n, d, device="cuda", dtype=dtype) * 2
+    m = torch.randn(n, d, device="cuda", dtype=dtype) * 8
+    go = torch.randn(n, d, device="cuda", dtype=dtype)
+    offsets = make_offsets(n, num_experts)
+    a = torch.full((num_experts + EXPERT_OFFSET, 3), 1 / 3, device="cuda")
+    w = torch.sigmoid(a)
+    s = (torch.rand(n, device="cuda") * 0.5 + 0.5) if with_scores else None
+    consts = (DEFAULT_GATE_BETA, DEFAULT_UP_BETA, *DEFAULT_SCALES, 1.0e6)
+    fwd_bytes = 3 * n * d * esz + (n * 4 if with_scores else 0)
+    bwd_bytes = 5 * n * d * esz + (n * 8 if with_scores else 0)
+
+    res = {}
+    # fused kernels, called directly (no autograd bookkeeping)
+    res["fused fwd"] = (bench(lambda: _ops.epolynorm_tanh_forward(
+        x, m, w, offsets, EXPERT_OFFSET, s, *consts)), fwd_bytes)
+    res["fused bwd"] = (bench(lambda: _ops.epolynorm_tanh_backward(
+        go, x, m, w, offsets, EXPERT_OFFSET, s, *consts)), bwd_bytes)
+
+    # eager reference: forward, and forward+backward via autograd
+    xr = x.clone().requires_grad_(True)
+    mr = m.clone().requires_grad_(True)
+    sr = s.reshape(-1, 1).clone().requires_grad_(True) if with_scores else None
+    wr = w.clone().requires_grad_(True)
+
+    def eager_fwd():
+        return grouped_epolynorm_tanh_ref(x,
+                                          m,
+                                          w,
+                                          offsets,
+                                          EXPERT_OFFSET,
+                                          scores=sr,
+                                          hidden_clamp=1.0e6)
+
+    def eager_fb():
+        out = grouped_epolynorm_tanh_ref(xr,
+                                         mr,
+                                         wr,
+                                         offsets,
+                                         EXPERT_OFFSET,
+                                         scores=sr,
+                                         hidden_clamp=1.0e6)
+        out.backward(go)
+
+    res["eager fwd"] = (bench(eager_fwd, iters=20), fwd_bytes)
+    t_fb = bench(eager_fb, iters=20)
+    res["eager bwd (fwd+bwd - fwd)"] = (t_fb - res["eager fwd"][0], bwd_bytes)
+
+    # compiled reference
+    try:
+        cref = torch.compile(grouped_epolynorm_tanh_ref, dynamic=False)
+
+        def comp_fwd():
+            return cref(x,
+                        m,
+                        w,
+                        offsets,
+                        EXPERT_OFFSET,
+                        scores=sr,
+                        hidden_clamp=1.0e6)
+
+        def comp_fb():
+            out = cref(xr,
+                       mr,
+                       wr,
+                       offsets,
+                       EXPERT_OFFSET,
+                       scores=sr,
+                       hidden_clamp=1.0e6)
+            out.backward(go)
+
+        res["compiled fwd"] = (bench(comp_fwd, iters=20), fwd_bytes)
+        t_fb = bench(comp_fb, iters=20)
+        res["compiled bwd (fwd+bwd - fwd)"] = (t_fb - res["compiled fwd"][0],
+                                               bwd_bytes)
+    except Exception as e:  # pragma: no cover
+        print(f"  (torch.compile failed: {e.__class__.__name__}: {e})")
+
+    # fused through autograd (what training uses), for the end-to-end number
+    xf = x.clone().requires_grad_(True)
+    mf = m.clone().requires_grad_(True)
+    sf = s.clone().requires_grad_(True) if with_scores else None
+    wf = w.clone().requires_grad_(True)
+
+    def fused_fb():
+        out = fused_grouped_epolynorm_tanh(xf,
+                                           mf,
+                                           wf,
+                                           offsets,
+                                           EXPERT_OFFSET,
+                                           scores=sf,
+                                           hidden_clamp=1.0e6)
+        out.backward(go)
+
+    res["fused fwd+bwd (autograd)"] = (bench(fused_fb), fwd_bytes + bwd_bytes)
+
+    tag = (f"N={n} D={d} E={num_experts} {str(dtype).split('.')[-1]} "
+           f"scores={with_scores}")
+    print(
+        f"\n== {tag}  (fwd min bytes {fwd_bytes/1e9:.3f} GB, bwd {bwd_bytes/1e9:.3f} GB)"
+    )
+    print(f"  {'variant':32s} {'ms':>9s} {'GB/s':>9s} {'% peak':>8s}")
+    for k, (ms, nb) in res.items():
+        gbs = nb / ms / 1e6
+        print(f"  {k:32s} {ms:9.3f} {gbs:9.0f} {100*gbs/peak:7.1f}%")
+    return res
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--d", type=int, nargs="+", default=[1152])
+    # 384B at local bs 12: 12*4096 tokens/rank x top_k 6 = 294912 routed rows per microbatch
+    ap.add_argument("--n", type=int, nargs="+", default=[294912, 65536, 8192])
+    # 384B: 384 experts / EP 8 = 48 local experts per rank
+    ap.add_argument("--experts", type=int, default=48)
+    ap.add_argument("--dtypes", nargs="+", default=["bfloat16", "float32"])
+    ap.add_argument("--no-scores", action="store_true")
+    args = ap.parse_args()
+
+    peak, how = theoretical_bw_gbs()
+    print(f"theoretical HBM bandwidth: {peak:.0f} GB/s  [{how}]")
+    for mb in (256, 2048):
+        c = copy_ceiling_gbs(mb * 2**20)
+        print(
+            f"device copy ceiling @ {mb} MiB: {c:.0f} GB/s ({100*c/peak:.1f}% of theoretical)"
+        )
+
+    for dt in args.dtypes:
+        dtype = getattr(torch, dt)
+        for d in args.d:
+            for n in args.n:
+                for with_scores in ((False, ) if args.no_scores else
+                                    (True, False)):
+                    run(n, d, args.experts, dtype, with_scores, peak)
+
+
+if __name__ == "__main__":
+    main()

build.toml

@@ -20,6 +20,7 @@ src = [
   "activation/rms_norm.cu",
   "activation/fused_add_rms_norm.cu",
   "activation/grouped_poly_norm.cu",
+  "activation/epolynorm_tanh.cu",
   "activation/cuda_compat.h",
   "activation/dispatch_utils.h",
   "activation/assert_utils.h",
@@ -35,6 +36,7 @@ src = [
   "activation/rms_norm.cu",
   "activation/fused_add_rms_norm.cu",
   "activation/grouped_poly_norm.cu",
+  "activation/epolynorm_tanh.cu",
   "activation/cuda_compat.h",
   "activation/dispatch_utils.h",
   "activation/assert_utils.h",

setup.py

@@ -24,6 +24,7 @@ CUDA_SOURCES = [
     "activation/rms_norm.cu",
     "activation/fused_add_rms_norm.cu",
     "activation/grouped_poly_norm.cu",
+    "activation/epolynorm_tanh.cu",
 ]
 
 CPP_SOURCES = [
@@ -58,6 +59,9 @@ NVCC_FLAGS = [
 cuda_version = tuple(int(x) for x in torch.version.cuda.split(".")[:2])
 if cuda_version >= (12, 8):
     NVCC_FLAGS.append("-gencode=arch=compute_100,code=sm_100")
+# B300 (sm_103) needs CUDA 12.9+ to be named explicitly; sm_100 SASS also runs there.
+if cuda_version >= (12, 9):
+    NVCC_FLAGS.append("-gencode=arch=compute_103,code=sm_103")
 
 CXX_FLAGS = ["-O3", CXX_STD]
 

tests/test_epolynorm_tanh.py

@@ -0,0 +1,348 @@
+"""Fused grouped EPolyNormTanh vs the eager reference: forward, backward
+(input, mul, per-expert weight, scores), expert_offset, padding rows,
+hidden_clamp masks, odd D (scalar path), the single-expert wrapper, opcheck."""
+
+import pytest
+import torch
+from activation.epolynorm_tanh import (DEFAULT_GATE_BETA, DEFAULT_SCALES,
+                                       DEFAULT_UP_BETA, _has_cuda_ops,
+                                       epolynorm_tanh_ref,
+                                       grouped_epolynorm_tanh_ref)
+
+if _has_cuda_ops:
+    from activation.epolynorm_tanh import (fused_epolynorm_tanh,
+                                           fused_grouped_epolynorm_tanh)
+
+from .utils import assert_close
+
+DTYPES = [torch.float, torch.bfloat16, torch.float16]
+NUM_TOKENS = [1, 1024, 4096]
+D = [1152, 256, 1280]
+NUM_EXPERTS = [8, 384]
+EXPERT_OFFSETS = [0, 4]
+CUDA_DEVICES = ["cuda:0"]
+
+pytestmark = pytest.mark.skipif(not _has_cuda_ops,
+                                reason="CUDA ops not available")
+
+
+def _tols(dtype):
+    # Both sides compute elementwise in fp32; the kernel is built with
+    # --use_fast_math (approximate division / exp inside tanhf), so fp32 agrees
+    # to ~1e-4 relative rather than to the last ulp. 16-bit outputs are
+    # dominated by their own rounding (bf16 eps = 7.8e-3).
+    if dtype == torch.float32:
+        return dict(atol=1e-3, rtol=1e-3)
+    return dict(atol=2e-2, rtol=2e-2)
+
+
+def _wgrad_tol(dtype):
+    return 5e-3 if dtype == torch.float32 else 2e-2
+
+
+def _make_inputs(n,
+                 d,
+                 num_experts,
+                 dtype,
+                 device,
+                 seed=0,
+                 scale=2.0,
+                 expert_offset=0,
+                 padding=0):
+    """Rows grouped by expert (random counts), optional padding tail."""
+    torch.manual_seed(seed)
+    probs = torch.ones(num_experts) / num_experts
+    assign = torch.multinomial(probs, n, replacement=True)
+    counts = torch.bincount(assign, minlength=num_experts)
+    offsets = torch.cumsum(counts, 0).to(torch.int32).to(device)
+    total = n + padding
+    x = torch.randn(total, d, device=device, dtype=dtype) * scale
+    m = torch.randn(total, d, device=device, dtype=dtype) * scale * 4
+    a = torch.full((num_experts + expert_offset, 3),
+                   1.0 / 3.0,
+                   device=device,
+                   dtype=torch.float32)
+    a = a + torch.randn_like(a) * 0.3
+    return x, m, a, offsets
+
+
+def _make_scores(n, device):
+    return torch.rand(n, 1, device=device, dtype=torch.float32) * 0.5 + 0.5
+
+
+def _run(fn,
+         x,
+         m,
+         a,
+         offsets,
+         expert_offset=0,
+         scores=None,
+         hidden_clamp=None):
+    xx = x.clone().detach().requires_grad_(True)
+    mm = m.clone().detach().requires_grad_(True)
+    aa = a.clone().detach().requires_grad_(True)
+    ss = scores.clone().detach().requires_grad_(
+        True) if scores is not None else None
+    out = fn(xx,
+             mm,
+             torch.sigmoid(aa),
+             offsets,
+             expert_offset,
+             scores=ss,
+             hidden_clamp=hidden_clamp)
+    # a random upstream gradient (not .sum()) so every element's chain is tested
+    torch.manual_seed(1)
+    go = torch.randn_like(out)
+    out.backward(go)
+    return out, xx.grad, mm.grad, aa.grad, (ss.grad
+                                            if ss is not None else None)
+
+
+def _check_wgrad(ga_c, ga_r, dtype, expert_offset=0):
+    """Per-expert [E, 3] weight grads: reductions over each expert's rows
+    (atomicAdd order differs), compared relative to the largest entry."""
+    assert ga_c.shape == ga_r.shape
+    if expert_offset > 0:  # rows below expert_offset get no gradient
+        assert (ga_c[:expert_offset] == 0).all()
+        assert (ga_r[:expert_offset] == 0).all()
+    rel = (ga_c - ga_r).abs().max() / ga_r.abs().max().clamp_min(1e-12)
+    assert rel < _wgrad_tol(dtype), (rel, ga_c, ga_r)
+
+
+@pytest.mark.parametrize("n", NUM_TOKENS)
+@pytest.mark.parametrize("d", D)
+@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_forward(n, d, num_experts, dtype, device):
+    torch.set_default_device(device)
+    x, m, a, offsets = _make_inputs(n, d, num_experts, dtype, device)
+    w = torch.sigmoid(a)
+    ref = grouped_epolynorm_tanh_ref(x, m, w, offsets)
+    out = fused_grouped_epolynorm_tanh(x, m, w, offsets)
+    assert out.shape == ref.shape and out.dtype == dtype
+    assert_close(out, ref, **_tols(dtype))
+
+
+@pytest.mark.parametrize("n", NUM_TOKENS)
+@pytest.mark.parametrize("d", D)
+@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("with_scores", [False, True])
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_backward(n, d, num_experts, dtype, with_scores, device):
+    torch.set_default_device(device)
+    x, m, a, offsets = _make_inputs(n, d, num_experts, dtype, device)
+    s = _make_scores(n, device) if with_scores else None
+    r = _run(grouped_epolynorm_tanh_ref, x, m, a, offsets, scores=s)
+    c = _run(fused_grouped_epolynorm_tanh, x, m, a, offsets, scores=s)
+    out_r, gx_r, gm_r, ga_r, gs_r = r
+    out_c, gx_c, gm_c, ga_c, gs_c = c
+    t = _tols(dtype)
+    assert_close(out_c, out_r, **t)
+    assert_close(gx_c, gx_r, **t)
+    assert_close(gm_c, gm_r, **t)
+    _check_wgrad(ga_c, ga_r, dtype)
+    if with_scores:
+        assert gs_c.shape == gs_r.shape
+        assert_close(gs_c, gs_r, atol=1e-3 * d**0.5, rtol=1e-2)
+
+
+@pytest.mark.parametrize("expert_offset", EXPERT_OFFSETS)
+@pytest.mark.parametrize("padding", [0, 37])
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_expert_offset_and_padding(expert_offset, padding, dtype, device):
+    """EP rank offset into the weight rows, and padding rows past offsets[-1]
+    (zero output, zero grads, no weight-grad contribution). The padding rows
+    hold NaN to prove they are never read into anything that matters."""
+    torch.set_default_device(device)
+    n, d, num_experts = 2048, 1152, 16
+    x, m, a, offsets = _make_inputs(n,
+                                    d,
+                                    num_experts,
+                                    dtype,
+                                    device,
+                                    expert_offset=expert_offset,
+                                    padding=padding)
+    if padding:
+        x[n:] = float("nan")
+        m[n:] = float("nan")
+    s = _make_scores(n + padding, device)
+    r = _run(grouped_epolynorm_tanh_ref, x, m, a, offsets, expert_offset, s)
+    c = _run(fused_grouped_epolynorm_tanh, x, m, a, offsets, expert_offset, s)
+    out_r, gx_r, gm_r, ga_r, gs_r = r
+    out_c, gx_c, gm_c, ga_c, gs_c = c
+    t = _tols(dtype)
+    assert_close(out_c[:n], out_r[:n], **t)
+    assert_close(gx_c[:n], gx_r[:n], **t)
+    assert_close(gm_c[:n], gm_r[:n], **t)
+    assert_close(gs_c[:n], gs_r[:n], atol=1e-3 * d**0.5, rtol=1e-2)
+    _check_wgrad(ga_c, ga_r, dtype, expert_offset)
+    if padding:
+        assert (out_c[n:] == 0).all()
+        assert (gx_c[n:] == 0).all()
+        assert (gm_c[n:] == 0).all()
+        assert (gs_c[n:] == 0).all()
+    assert torch.isfinite(ga_c).all()
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("hidden_clamp", [1000000.0, 3.0, 0.5])
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_hidden_clamp(dtype, hidden_clamp, device):
+    """clamp_max on x (upper only) and symmetric clamp on mul, with the right
+    zero-gradient masks."""
+    torch.set_default_device(device)
+    x, m, a, offsets = _make_inputs(4096, 1152, 8, dtype, device, scale=1.0)
+    s = _make_scores(4096, device)
+    r = _run(grouped_epolynorm_tanh_ref, x, m, a, offsets, 0, s, hidden_clamp)
+    c = _run(fused_grouped_epolynorm_tanh, x, m, a, offsets, 0, s,
+             hidden_clamp)
+    out_r, gx_r, gm_r, ga_r, gs_r = r
+    out_c, gx_c, gm_c, ga_c, gs_c = c
+    t = _tols(dtype)
+    assert_close(out_c, out_r, **t)
+    assert_close(gx_c, gx_r, **t)
+    assert_close(gm_c, gm_r, **t)
+    if hidden_clamp < 10:
+        assert (gx_c[x > hidden_clamp] == 0).all()
+        assert (gm_c[m.abs() > hidden_clamp] == 0).all()
+        assert (gx_c[x < -hidden_clamp] != 0).any()  # lower side NOT clamped
+    _check_wgrad(ga_c, ga_r, dtype)
+    assert_close(gs_c, gs_r, atol=1e-3 * 1152**0.5, rtol=1e-2)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_single_expert_wrapper(dtype, device):
+    """fused_epolynorm_tanh (dense FFN: weight [3], any leading shape) ==
+    the single-expert eager reference."""
+    torch.set_default_device(device)
+    x, m, a, _ = _make_inputs(2048, 1152, 1, dtype, device, seed=2)
+    x3, m3 = x.reshape(4, 512, 1152), m.reshape(4, 512, 1152)
+    a1 = a.reshape(3)
+    xr, mr, ar = (t.clone().requires_grad_(True) for t in (x3, m3, a1))
+    xc, mc, ac = (t.clone().requires_grad_(True) for t in (x3, m3, a1))
+    out_r = epolynorm_tanh_ref(xr, mr, torch.sigmoid(ar), hidden_clamp=1e6)
+    out_c = fused_epolynorm_tanh(xc, mc, torch.sigmoid(ac), hidden_clamp=1e6)
+    assert out_c.shape == x3.shape
+    t = _tols(dtype)
+    assert_close(out_c, out_r, **t)
+    torch.manual_seed(4)
+    go = torch.randn_like(out_r)
+    out_r.backward(go)
+    out_c.backward(go)
+    assert_close(xc.grad, xr.grad, **t)
+    assert_close(mc.grad, mr.grad, **t)
+    _check_wgrad(ac.grad.reshape(1, 3), ar.grad.reshape(1, 3), dtype)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_module_matches_eager_module(dtype, device):
+    """activation.layers.EPolyNormTanh (per-expert weight [E, 3]) == an eager
+    per-expert module written from the activation's definition."""
+    import math
+
+    import torch.nn as nn
+    from activation.layers import EPolyNormTanh
+
+    class Eager(nn.Module):
+
+        def __init__(self, num_experts):
+            super().__init__()
+            self.weight = nn.Parameter(
+                torch.full((num_experts, 3), 1 / 3, dtype=torch.float32))
+
+        def forward(self,
+                    x,
+                    mul,
+                    offsets,
+                    expert_offset=0,
+                    scores=None,
+                    hidden_clamp=None):
+            orig = x.dtype
+            n = int(offsets[-1].item())
+            idx = torch.bucketize(torch.arange(n, device=x.device),
+                                  offsets,
+                                  right=True) + expert_offset
+            w = torch.sigmoid(self.weight.float())[idx]
+            x, mul = x[:n], mul[:n]
+            if hidden_clamp is not None:
+                x = x.clamp(max=hidden_clamp)
+                mul = mul.clamp(-hidden_clamp, hidden_clamp)
+            g, u = x.float(), mul.float()
+            s0, s1, s2 = 1.0, math.sqrt(3.0), math.sqrt(15.0)
+            q0 = torch.tanh(g / s0)
+            q1 = torch.tanh((g * g) / s1)
+            q2 = torch.tanh((g * g * g) / s2)
+            poly = 4.0 * torch.tanh(
+                (w[:, 0:1] * q0 + w[:, 1:2] * q1 + w[:, 2:3] * q2) / 4.0)
+            r = poly * (25.0 * torch.tanh(u / 25.0))
+            if scores is not None:
+                r = r * scores[:n]
+            return r.to(orig)
+
+    torch.set_default_device(device)
+    E, eo = 12, 4
+    x, m, _, offsets = _make_inputs(2048,
+                                    1152,
+                                    E,
+                                    dtype,
+                                    device,
+                                    seed=3,
+                                    expert_offset=eo)
+    s = _make_scores(2048, device)
+    eager, fused = Eager(E + eo), EPolyNormTanh(E + eo)
+    with torch.no_grad():
+        torch.manual_seed(7)
+        eager.weight.add_(torch.randn_like(eager.weight) * 0.2)
+        fused.weight.copy_(eager.weight)
+    xe, me, se = (t.clone().requires_grad_(True) for t in (x, m, s))
+    xf, mf, sf = (t.clone().requires_grad_(True) for t in (x, m, s))
+    oe = eager(xe, me, offsets, eo, scores=se, hidden_clamp=1e6)
+    of = fused(xf, mf, offsets, eo, scores=sf, hidden_clamp=1e6)
+    t = _tols(dtype)
+    assert_close(of, oe, **t)
+    torch.manual_seed(5)
+    go = torch.randn_like(oe)
+    oe.backward(go)
+    of.backward(go)
+    assert_close(xf.grad, xe.grad, **t)
+    assert_close(mf.grad, me.grad, **t)
+    assert_close(sf.grad, se.grad, atol=1e-3 * 1152**0.5, rtol=1e-2)
+    _check_wgrad(fused.weight.grad, eager.weight.grad, dtype, eo)
+    # every local expert actually received a gradient
+    assert (fused.weight.grad[eo:].abs().sum(dim=1) > 0).all()
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_no_nan_inf_large_inputs(dtype, device):
+    torch.set_default_device(device)
+    x, m, a, offsets = _make_inputs(1024, 1152, 8, dtype, device, scale=50.0)
+    out, gx, gm, ga, _ = _run(fused_grouped_epolynorm_tanh, x, m, a, offsets)
+    for name, t in (("out", out), ("gx", gx), ("gm", gm), ("ga", ga)):
+        assert torch.isfinite(t).all(), name
+
+
+@pytest.mark.parametrize("with_scores", [False, True])
+@pytest.mark.parametrize("device", CUDA_DEVICES)
+def test_opcheck(with_scores, device):
+    torch.set_default_device(device)
+    x, m, a, offsets = _make_inputs(64, 256, 4, torch.float32, device)
+    w = torch.sigmoid(a)
+    s = _make_scores(64, device).reshape(-1) if with_scores else None
+    g0, g1, s0, s1, s2 = (DEFAULT_GATE_BETA, DEFAULT_UP_BETA, *DEFAULT_SCALES)
+    torch.library.opcheck(torch.ops._activation.epolynorm_tanh_forward,
+                          (x, m, w, offsets, 0, s, g0, g1, s0, s1, s2, -1.0))
+    go = torch.randn_like(x)
+    # The backward accumulates the per-expert weight gradient with atomicAdd
+    # (summation order differs run to run), so the eager-vs-AOT output
+    # comparison is skipped for it; schema / fake-tensor / autograd checks run.
+    torch.library.opcheck(
+        torch.ops._activation.epolynorm_tanh_backward,
+        (go, x, m, w, offsets, 0, s, g0, g1, s0, s1, s2, -1.0),
+        test_utils=("test_schema", "test_autograd_registration",
+                    "test_faketensor"))

torch-ext/activation/__init__.py

@@ -2,6 +2,9 @@ import torch
 
 from . import layers, parallel_style
 from ._ops import ops
+from .epolynorm_tanh import (epolynorm_tanh_ref, fused_epolynorm_tanh,
+                             fused_grouped_epolynorm_tanh,
+                             grouped_epolynorm_tanh_ref)
 from .fused_rope import fused_kv_split_rope_cat, fused_q_rope_inplace
 from .grouped_poly_norm import fused_mul_grouped_poly_norm
 from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
@@ -49,6 +52,10 @@ __all__ = [
     "poly_norm",
     "fused_mul_poly_norm",
     "fused_mul_grouped_poly_norm",
+    "fused_epolynorm_tanh",
+    "fused_grouped_epolynorm_tanh",
+    "grouped_epolynorm_tanh_ref",
+    "epolynorm_tanh_ref",
     "rms_norm",
     "fused_add_rms_norm",
     "fused_q_rope_inplace",

torch-ext/activation/epolynorm_tanh.py

@@ -0,0 +1,269 @@
+"""Grouped (per-expert) EPolyNormTanh — CUDA kernel with autograd wrapper and a
+PyTorch reference.
+
+Tanh-bounded elementwise PolyNorm with SiTU-style caps (the ``epolynorm_tanh``
+MoE/FFN activation of llm-training), with one coefficient triple per expert
+like the grouped PolyNorm:
+
+    e      = expert of the row (rows grouped by expert; offsets = cumsum of
+             tokens per local expert; coefficient row = e + expert_offset)
+    g      = clamp_max(x, C)                       (C = hidden_clamp, None = off)
+    u      = clamp(mul, -C, C)
+    q_k    = tanh(g^(k+1) / s_k)                   k = 0, 1, 2
+    poly   = gate_beta * tanh((w_e0 q_0 + w_e1 q_1 + w_e2 q_2) / gate_beta)
+    up     = up_beta * tanh(u / up_beta)
+    out    = poly * up * score                     (score: optional per-row fp32)
+
+``weight`` holds the coefficients *after* the sigmoid (fp32 [num_total_experts,
+3]); the caller keeps the raw parameters and lets autograd carry the gradient
+back through ``torch.sigmoid``. The kernel computes in fp32 and returns the
+input dtype, exactly like the eager implementation. Rows at or beyond
+``offsets[-1]`` are padding: zero output and zero gradients.
+
+CUDA kernel (activation/epolynorm_tanh.cu):
+  - one warp per row, persistent grid, 128-bit vectorized loads (width 8 for
+    16-bit types, 4 for fp32), scalar fallback otherwise
+  - expert lookup by in-kernel binary search over ``offsets`` (per row range)
+  - forward: one pass, no saved intermediates (backward recomputes the tanh's)
+  - backward: one pass -> grad_input, grad_mul, grad_weight (fp32
+    [num_total_experts, 3], warp-reduced then atomicAdd per expert) and
+    grad_scores (per-row warp reduction)
+"""
+
+import math
+
+import torch
+from torch import Tensor
+
+_ops = None
+try:
+    import _activation  # noqa: F401 — local setup.py build
+    _ops = torch.ops._activation
+except Exception:
+    try:
+        from ._ops import ops as _ops
+    except Exception:
+        pass
+
+_has_cuda_ops = _ops is not None and hasattr(_ops, "epolynorm_tanh_forward")
+
+DEFAULT_GATE_BETA = 4.0
+DEFAULT_UP_BETA = 25.0
+DEFAULT_SCALES = (1.0, math.sqrt(3.0), math.sqrt(15.0))
+
+if _has_cuda_ops:
+    try:
+
+        @torch.library.register_fake("_activation::epolynorm_tanh_forward")
+        def _fwd_fake(input, mul, weight, offsets, expert_offset, scores,
+                      gate_beta, up_beta, s0, s1, s2, hidden_clamp):
+            return torch.empty_like(input)
+
+        @torch.library.register_fake("_activation::epolynorm_tanh_backward")
+        def _bwd_fake(grad_output, input, mul, weight, offsets, expert_offset,
+                      scores, gate_beta, up_beta, s0, s1, s2, hidden_clamp):
+            gs = (torch.empty(
+                input.shape[0], 1, dtype=torch.float32, device=input.device)
+                  if scores is not None else torch.empty(
+                      0, dtype=torch.float32, device=input.device))
+            return (torch.empty_like(input), torch.empty_like(mul),
+                    torch.empty_like(weight), gs)
+    except Exception:
+        pass  # already registered
+
+
+# ---------------------------------------------------------------------------
+# PyTorch reference (mirrors llm-training GroupedExpertsEPolyNormTanh)
+# ---------------------------------------------------------------------------
+def grouped_epolynorm_tanh_ref(
+    input: Tensor,
+    mul: Tensor,
+    weight: Tensor,
+    offsets: Tensor,
+    expert_offset: int = 0,
+    scores: Tensor | None = None,
+    hidden_clamp: float | None = None,
+    gate_beta: float = DEFAULT_GATE_BETA,
+    up_beta: float = DEFAULT_UP_BETA,
+    scales: tuple[float, float, float] = DEFAULT_SCALES,
+) -> Tensor:
+    """Eager reference. ``weight`` = sigmoid(a) as fp32 [num_total_experts, 3].
+
+    Rows are mapped to experts with ``torch.bucketize`` over ``offsets``; rows
+    at or beyond ``offsets[-1]`` are padding and come back as zeros.
+    """
+    orig_dtype = input.dtype
+    num_valid = int(offsets[-1].item())
+    positions = torch.arange(num_valid, device=input.device)
+    expert_idx = torch.bucketize(positions, offsets,
+                                 right=True) + expert_offset
+    w = weight.float()[expert_idx]  # [num_valid, 3]
+
+    x = input[:num_valid]
+    m = mul[:num_valid]
+    if hidden_clamp is not None:
+        x = x.clamp(max=hidden_clamp)
+        m = m.clamp(-hidden_clamp, hidden_clamp)
+    g = x.float()
+    u = m.float()
+    s0, s1, s2 = scales
+    q0 = torch.tanh(g / s0)
+    q1 = torch.tanh((g * g) / s1)
+    q2 = torch.tanh((g * g * g) / s2)
+    poly = gate_beta * torch.tanh(
+        (w[:, 0:1] * q0 + w[:, 1:2] * q1 + w[:, 2:3] * q2) / gate_beta)
+    up_cap = up_beta * torch.tanh(u / up_beta)
+    result = poly * up_cap
+    if scores is not None:
+        assert scores.dtype == torch.float32, (
+            f"scores dtype must be float32, got {scores.dtype}")
+        result = result * scores[:num_valid].reshape(-1, 1)
+    result = result.to(orig_dtype)
+    if input.shape[0] > num_valid:
+        pad = torch.zeros(input.shape[0] - num_valid,
+                          input.shape[-1],
+                          dtype=orig_dtype,
+                          device=input.device)
+        result = torch.cat([result, pad], dim=0)
+    return result
+
+
+def epolynorm_tanh_ref(
+    input: Tensor,
+    mul: Tensor,
+    weight: Tensor,
+    scores: Tensor | None = None,
+    hidden_clamp: float | None = None,
+    gate_beta: float = DEFAULT_GATE_BETA,
+    up_beta: float = DEFAULT_UP_BETA,
+    scales: tuple[float, float, float] = DEFAULT_SCALES,
+) -> Tensor:
+    """Single-expert eager reference (dense FFN). ``weight`` = fp32 [3]."""
+    D = input.shape[-1]
+    x = input.reshape(-1, D)
+    offsets = torch.tensor([x.shape[0]], dtype=torch.int32, device=x.device)
+    out = grouped_epolynorm_tanh_ref(x, mul.reshape(-1, D),
+                                     weight.reshape(1, 3), offsets, 0, scores,
+                                     hidden_clamp, gate_beta, up_beta, scales)
+    return out.reshape(input.shape)
+
+
+# ---------------------------------------------------------------------------
+# CUDA autograd function
+# ---------------------------------------------------------------------------
+if _has_cuda_ops:
+
+    class _GroupedEPolyNormTanhFn(torch.autograd.Function):
+
+        @staticmethod
+        def forward(input, mul, weight, offsets, expert_offset, scores,
+                    gate_beta, up_beta, s0, s1, s2, hidden_clamp):
+            return _ops.epolynorm_tanh_forward(input, mul, weight, offsets,
+                                               expert_offset, scores,
+                                               gate_beta, up_beta, s0, s1, s2,
+                                               hidden_clamp)
+
+        @staticmethod
+        def setup_context(ctx, inputs, output):
+            (input, mul, weight, offsets, expert_offset, scores, gate_beta,
+             up_beta, s0, s1, s2, hidden_clamp) = inputs
+            ctx.save_for_backward(input, mul, weight, offsets, scores)
+            ctx.expert_offset = expert_offset
+            ctx.consts = (gate_beta, up_beta, s0, s1, s2, hidden_clamp)
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            input, mul, weight, offsets, scores = ctx.saved_tensors
+            gi, gm, gw, gs = _ops.epolynorm_tanh_backward(
+                grad_output.contiguous(), input, mul, weight, offsets,
+                ctx.expert_offset, scores, *ctx.consts)
+            if scores is None:
+                gs = None
+            elif scores.dim() == 1:
+                gs = gs.reshape(-1)
+            return (gi, gm, gw, None, None, gs, None, None, None, None, None,
+                    None)
+
+    def fused_grouped_epolynorm_tanh(
+        input: Tensor,
+        mul: Tensor,
+        weight: Tensor,
+        offsets: Tensor,
+        expert_offset: int = 0,
+        scores: Tensor | None = None,
+        hidden_clamp: float | None = None,
+        gate_beta: float = DEFAULT_GATE_BETA,
+        up_beta: float = DEFAULT_UP_BETA,
+        scales: tuple[float, float, float] = DEFAULT_SCALES,
+    ) -> Tensor:
+        """Fused grouped EPolyNormTanh.
+
+        Args:
+            input: (N, D) gate projection (w1 x), fp32/bf16/fp16; rows grouped
+                by expert
+            mul: (N, D) up projection (w3 x), same dtype
+            weight: (num_total_experts, 3) fp32 = sigmoid(a), same device
+            offsets: (num_experts,) int32 cumsum of tokens per local expert;
+                rows >= offsets[-1] are padding (zero output / zero grads)
+            expert_offset: global index of local expert 0 (EP rank offset)
+            scores: optional (N,) or (N, 1) fp32 routing scores
+            hidden_clamp: clamp_max(input, C) and clamp(mul, -C, C); None = off
+            gate_beta, up_beta, scales: the activation's constants
+        Returns:
+            (N, D) in the input dtype
+        """
+        assert input.dim() == 2 and mul.shape == input.shape, (input.shape,
+                                                               mul.shape)
+        x = input if input.is_contiguous() else input.contiguous()
+        m = mul if mul.is_contiguous() else mul.contiguous()
+        if offsets.dtype != torch.int32:
+            offsets = offsets.to(torch.int32)
+        offsets = offsets.contiguous()
+        if weight.dtype != torch.float32:
+            weight = weight.float()
+        weight = weight.contiguous()
+        if scores is not None:
+            assert scores.dtype == torch.float32, (
+                f"scores dtype must be float32, got {scores.dtype}")
+            scores = scores.contiguous()
+        clamp_val = -1.0 if hidden_clamp is None else float(hidden_clamp)
+        s0, s1, s2 = (float(v) for v in scales)
+        return _GroupedEPolyNormTanhFn.apply(x, m, weight, offsets,
+                                             int(expert_offset), scores,
+                                             float(gate_beta), float(up_beta),
+                                             s0, s1, s2, clamp_val)
+
+    def fused_epolynorm_tanh(
+        input: Tensor,
+        mul: Tensor,
+        weight: Tensor,
+        scores: Tensor | None = None,
+        hidden_clamp: float | None = None,
+        gate_beta: float = DEFAULT_GATE_BETA,
+        up_beta: float = DEFAULT_UP_BETA,
+        scales: tuple[float, float, float] = DEFAULT_SCALES,
+    ) -> Tensor:
+        """Single-expert fused EPolyNormTanh (dense FFN): ``weight`` fp32 [3],
+        any leading shape."""
+        D = input.shape[-1]
+        x = input.reshape(-1, D)
+        offsets = torch.tensor([x.shape[0]],
+                               dtype=torch.int32,
+                               device=x.device)
+        out = fused_grouped_epolynorm_tanh(x, mul.reshape(-1, D),
+                                           weight.reshape(1, 3), offsets, 0,
+                                           scores, hidden_clamp, gate_beta,
+                                           up_beta, scales)
+        return out.reshape(input.shape)
+
+else:
+
+    def fused_grouped_epolynorm_tanh(*args, **kwargs):
+        raise RuntimeError(
+            "CUDA ops not available. Build with setup.py or kernel-builder "
+            "to use fused_grouped_epolynorm_tanh.")
+
+    def fused_epolynorm_tanh(*args, **kwargs):
+        raise RuntimeError(
+            "CUDA ops not available. Build with setup.py or kernel-builder "
+            "to use fused_epolynorm_tanh.")

torch-ext/activation/layers.py

@@ -93,3 +93,78 @@ class FusedAddRMSNorm(nn.Module):
         Resets parameters based on their initialization used in __init__.
         """
         init.ones_(self.weight)
+
+
+class EPolyNormTanh(nn.Module):
+    """Fused grouped (per-expert) tanh-bounded elementwise PolyNorm with
+    SiTU-style caps.
+
+    ``out = gate_beta * tanh(sum_k sigmoid(a_ek) tanh(x^(k+1)/s_k) / gate_beta)
+           * up_beta * tanh(mul / up_beta) [* scores]``
+
+    ``weight`` is a fp32 [num_experts, 3] parameter (the raw coefficients
+    a_ek, passed through a sigmoid), one triple per expert like the grouped
+    PolyNorm. Drop-in for llm-training's ``GroupedExpertsEPolyNormTanh``: same
+    ``forward(x, mul, offsets=None, expert_offset=0, scores=None,
+    hidden_clamp=None)`` signature. With ``offsets=None`` (dense FFN) all rows
+    use expert 0.
+    """
+
+    def __init__(
+            self,
+            num_experts: int = 1,
+            gate_beta: float = 4.0,
+            up_beta: float = 25.0,
+            scales=(1.0, 3.0**0.5, 15.0**0.5),
+            a_init: float = 1.0 / 3.0,
+    ):
+        super().__init__()
+        from .epolynorm_tanh import fused_grouped_epolynorm_tanh
+        if gate_beta <= 0.0 or up_beta <= 0.0:
+            raise ValueError("gate_beta and up_beta must be positive.")
+        self._fn = fused_grouped_epolynorm_tanh
+        self.num_experts = int(num_experts)
+        self.gate_beta = float(gate_beta)
+        self.up_beta = float(up_beta)
+        self._scales = tuple(float(s) for s in scales)
+        self._a_init = float(a_init)
+        self.weight = nn.Parameter(
+            torch.full((self.num_experts, 3),
+                       self._a_init,
+                       dtype=torch.float32))
+
+    def reset_parameters(self) -> None:
+        with torch.no_grad():
+            init.constant_(self.weight, self._a_init)
+
+    def sigmoid_weight(self) -> torch.Tensor:
+        return torch.sigmoid(self.weight.float())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mul: torch.Tensor,
+        offsets: torch.Tensor | None = None,
+        expert_offset: int = 0,
+        scores: torch.Tensor | None = None,
+        hidden_clamp: float | None = None,
+    ) -> torch.Tensor:
+        shape = x.shape
+        D = shape[-1]
+        x = x.reshape(-1, D)
+        mul = mul.reshape(-1, D)
+        if offsets is None:
+            offsets = torch.tensor([x.shape[0]],
+                                   dtype=torch.int32,
+                                   device=x.device)
+        out = self._fn(x,
+                       mul,
+                       self.sigmoid_weight(),
+                       offsets,
+                       expert_offset,
+                       scores=scores,
+                       hidden_clamp=hidden_clamp,
+                       gate_beta=self.gate_beta,
+                       up_beta=self.up_beta,
+                       scales=self._scales)
+        return out.reshape(shape)

torch-ext/torch_binding.cpp

@@ -83,6 +83,20 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
           "Tensor, Tensor, Tensor, Tensor)");
   ops.impl("grouped_poly_norm_backward_scored", torch::kCUDA,
            &grouped_poly_norm_backward_scored);
+
+  // epolynorm_tanh (grouped, per-expert coefficients; scores optional,
+  // hidden_clamp < 0 = off)
+  ops.def("epolynorm_tanh_forward(Tensor input, Tensor mul, Tensor weight, "
+          "Tensor offsets, int expert_offset, Tensor? scores, float gate_beta, "
+          "float up_beta, float s0, float s1, float s2, float hidden_clamp) "
+          "-> Tensor");
+  ops.impl("epolynorm_tanh_forward", torch::kCUDA, &epolynorm_tanh_forward);
+
+  ops.def("epolynorm_tanh_backward(Tensor grad_output, Tensor input, Tensor "
+          "mul, Tensor weight, Tensor offsets, int expert_offset, "
+          "Tensor? scores, float gate_beta, float up_beta, float s0, float s1, "
+          "float s2, float hidden_clamp) -> (Tensor, Tensor, Tensor, Tensor)");
+  ops.impl("epolynorm_tanh_backward", torch::kCUDA, &epolynorm_tanh_backward);
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -2,6 +2,9 @@
 
 #include <torch/torch.h>
 
+#include <optional>
+#include <tuple>
+
 void poly_norm(torch::Tensor &out, const torch::Tensor &input,
                const torch::Tensor &weights, const torch::Tensor &bias,
                double eps);
@@ -69,3 +72,27 @@ grouped_poly_norm_backward_scored(
     const torch::Tensor &bias, const torch::Tensor &offsets,
     const torch::Tensor &inv_rms, const torch::Tensor &scores, double eps,
     int64_t expert_offset, double hidden_clamp);
+
+// Grouped EPolyNormTanh: tanh-bounded elementwise PolyNorm with SiTU-style
+// caps, per-expert coefficients. weight = sigmoid(a) fp32 [num_total_experts,
+// 3]; offsets int32 [num_experts] = cumsum of tokens per expert (rows >=
+// offsets[-1] are padding); the coefficient row is e + expert_offset; scores
+// optional fp32 [N]; hidden_clamp < 0 = disabled.
+torch::Tensor
+epolynorm_tanh_forward(const torch::Tensor &input, const torch::Tensor &mul,
+                       const torch::Tensor &weight,
+                       const torch::Tensor &offsets, int64_t expert_offset,
+                       const std::optional<torch::Tensor> &scores,
+                       double gate_beta, double up_beta, double s0, double s1,
+                       double s2, double hidden_clamp);
+
+// -> (grad_input, grad_mul, grad_weight fp32 [num_total_experts, 3],
+// grad_scores fp32 [N, 1], or an empty [0] tensor when scores is None)
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+epolynorm_tanh_backward(const torch::Tensor &grad_output,
+                        const torch::Tensor &input, const torch::Tensor &mul,
+                        const torch::Tensor &weight,
+                        const torch::Tensor &offsets, int64_t expert_offset,
+                        const std::optional<torch::Tensor> &scores,
+                        double gate_beta, double up_beta, double s0, double s1,
+                        double s2, double hidden_clamp);