Add Snake activation functor for EVT (#3184)

Introduces cutlass::epilogue::thread::Snake, a two-operand activation
functor implementing Snake_a(x) = x + (1/a) * sin^2(a*x) from
Ziyin et al. 2020 (arXiv:2006.08195). The per-channel learnable
frequency `a` flows through an EVT child (e.g. Sm90RowBroadcast),
composing into Sm90EVT<Sm90Compute<Snake, ...>, x_node, alpha_node>
for fused GEMM+Snake epilogues used in neural vocoders.

Adds unit tests in test/unit/epilogue/thread/activation.cu covering
f32 and bf16 paths, validated against float64 reference goldens.

Closes #3141
This commit is contained in:
Emre Albayrak
2026-05-12 05:09:53 +03:00
committed by GitHub
parent 1d9e1f6d7a
commit 2e56847d72
2 changed files with 367 additions and 1 deletions

View File

@@ -905,6 +905,41 @@ struct ElementwiseFilter<Array<T, N> > {
}
};
// Snake activation: Snake_a(x) = x + (1/a) * sin^2(a*x)
// Introduced in Ziyin, Hartwig, Ueda, "Neural Networks Fail to Learn
// Periodic Functions and How to Fix It," NeurIPS 2020 (arXiv:2006.08195).
// The per-channel learnable frequency `a` is passed as the second operand
// (intended to flow through an EVT child such as Sm90RowBroadcast).
// Caller must ensure a != 0 (the formula is singular at a = 0).
template <typename T>
struct Snake {
static const bool kIsHeavy = true;
CUTLASS_HOST_DEVICE
T operator()(T const& x, T const& alpha) const {
float xf = float(x);
float af = float(alpha);
float s = fast_sin(af * xf);
return T(xf + s * s / af);
}
};
template <typename T, int N>
struct Snake<Array<T, N>> {
static const bool kIsHeavy = true;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& x, Array<T, N> const& alpha) const {
Array<T, N> result;
Snake<T> scalar_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < N; ++i){
result[i] = scalar_op(x[i], alpha[i]);
}
return result;
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace thread