feat: Triton fused SoftSignGLU kernel for LYNXNet2

configs/acoustic.yaml

@@ -3,6 +3,9 @@ base_config:
 
 task_cls: training.acoustic_task.AcousticTask
 
+# Enable Triton-fused Linear+SoftSignGLU kernels for LYNXNet2 backbones.
+use_fused_kernels: false
+
 dictionaries: {}
 extra_phonemes: []
 merged_phoneme_groups: []

configs/variance.yaml

@@ -3,6 +3,9 @@ base_config:
 
 task_cls: training.variance_task.VarianceTask
 
+# Enable Triton-fused Linear+SoftSignGLU kernels for LYNXNet2 backbones.
+use_fused_kernels: false
+
 dictionaries: {}
 extra_phonemes: []
 merged_phoneme_groups: []

modules/backbones/lynxnet2.py

@@ -2,7 +2,9 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, ATanGLU, Transpose, AdamWLinear
+from modules.commons.common_layers import (
+    SinusoidalPosEmb, SwiGLU, ATanGLU, SoftSignGLU, Transpose, AdamWLinear
+)
 from utils.hparams import hparams
 
 
@@ -14,6 +16,8 @@ def __init__(self, dim, expansion_factor, kernel_size=31, dropout=0., glu_type='
             _glu = SwiGLU()
         elif glu_type == 'atanglu':
             _glu = ATanGLU()
+        elif glu_type == 'softsign_glu':
+            _glu = SoftSignGLU()
         else:
             raise ValueError(f'{glu_type} is not a valid activation')
         if float(dropout) > 0.:

modules/commons/common_layers.py

@@ -175,6 +175,47 @@ def forward(self, x):
             return out * torch.atan(gate)
 
 
+class SoftSignGLUFunction(torch.autograd.Function):
+    """ATanGLUFunction-style memory trick for SoftSignGLU.
+
+    softsign'(x) = 1/(1+|x|)^2 = (1-|softsign(x)|)^2, so both partial
+    derivatives of y = out * softsign(gate) are precomputable in forward:
+      dy/dout = softsign(gate)
+      dy/dgate = out * (1-|softsign(gate)|)^2
+    Saves 2 tensors (vs 3 for naive autograd) and backward is two pure
+    multiplies with no softsign recompute.
+    """
+    @staticmethod
+    def forward(ctx, out, gate):
+        ss_gate = torch.nn.functional.softsign(gate)
+        decay_out = out * (1.0 - ss_gate.abs()).square()
+        ctx.save_for_backward(ss_gate, decay_out)
+        return out * ss_gate
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        ss_gate, decay_out = ctx.saved_tensors
+        return grad_output * ss_gate, grad_output * decay_out
+
+
+class SoftSignGLU(nn.Module):
+    """Gated Linear Unit with SoftSign gate: out * softsign(gate).
+
+    More numerically stable than ATanGLU (no approximation needed in
+    Triton kernels) while providing similar gating behavior.
+    """
+    def __init__(self, dim=-1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        out, gate = torch.split(x, x.size(self.dim) // 2, dim=self.dim)
+        if self.training:
+            return SoftSignGLUFunction.apply(out, gate)
+        else:
+            return out * torch.nn.functional.softsign(gate)
+
+
 class AdamWConv1d(torch.nn.Conv1d):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)

modules/kernels/__init__.py

@@ -0,0 +1 @@
+# Fused kernels for LYNXNet2 optimization