feat: Triton fused SoftSignGLU kernel for LYNXNet2[WIP]

configs/acoustic.yaml

@@ -82,7 +82,7 @@ backbone_args:
   kernel_size: 31
   dropout_rate: 0.0
   use_conditioner_cache: true
-  glu_type: 'atanglu'
+  glu_type: 'softsign_glu'
 main_loss_type: l2
 main_loss_log_norm: false
 schedule_type: 'linear'

configs/base.yaml

@@ -82,6 +82,11 @@ pl_trainer_strategy:
   find_unused_parameters: false
 nccl_p2p: true
 
+###########
+# fusing
+###########
+use_fused_kernels: false
+
 ###########
 # finetune
 ###########

configs/templates/config_acoustic.yaml

@@ -90,7 +90,7 @@ backbone_args:
   kernel_size: 31
   dropout_rate: 0.0
   use_conditioner_cache: true
-  glu_type: 'atanglu'
+  glu_type: 'softsign_glu'
 shallow_diffusion_args:
   train_aux_decoder: true
   train_diffusion: true

configs/templates/config_duration.yaml

@@ -103,7 +103,7 @@ pitch_prediction_args:
     num_channels: 512
     dropout_rate: 0.0
     use_conditioner_cache: true
-    glu_type: 'atanglu'
+    glu_type: 'softsign_glu'
 
 variances_prediction_args:
   total_repeat_bins: 72
@@ -113,7 +113,7 @@ variances_prediction_args:
     num_channels: 384
     dropout_rate: 0.0
     use_conditioner_cache: true
-    glu_type: 'atanglu'
+    glu_type: 'softsign_glu'
 
 lambda_dur_loss: 1.0
 lambda_pitch_loss: 1.0

configs/templates/config_variance.yaml

@@ -103,7 +103,7 @@ pitch_prediction_args:
     num_channels: 512
     dropout_rate: 0.0
     use_conditioner_cache: true
-    glu_type: 'atanglu'
+    glu_type: 'softsign_glu'
 
 variances_prediction_args:
   total_repeat_bins: 72
@@ -113,7 +113,7 @@ variances_prediction_args:
     num_channels: 384
     dropout_rate: 0.0
     use_conditioner_cache: true
-    glu_type: 'atanglu'
+    glu_type: 'softsign_glu'
 
 lambda_dur_loss: 1.0
 lambda_pitch_loss: 1.0

configs/variance.yaml

@@ -74,7 +74,7 @@ pitch_prediction_args:
     num_channels: 512
     dropout_rate: 0.0
     use_conditioner_cache: true
-    glu_type: 'atanglu'
+    glu_type: 'softsign_glu'
 
 energy_db_min: -96.0
 energy_db_max: -12.0
@@ -99,7 +99,7 @@ variances_prediction_args:
     num_channels: 384
     dropout_rate: 0.0
     use_conditioner_cache: true
-    glu_type: 'atanglu'
+    glu_type: 'softsign_glu'
 
 lambda_dur_loss: 1.0
 lambda_pitch_loss: 1.0

modules/backbones/lynxnet2.py

@@ -2,7 +2,7 @@
 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 +14,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,21 @@ def forward(self, x):
             return out * torch.atan(gate)
 
 
+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)
+        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