fix: elem kernel extra N arg; variance nested glu_type

- Remove extra N arg in _softsign_glu_bwd_elem_kernel call (M,N,N → M,N)
- Read variance glu_type from nested predictor config paths
- Assert both predictors use softsign_glu before patching

Author: KakaruHayate

modules/kernels/fused_linear_softsign_glu.py

@@ -337,7 +337,7 @@ def elem_grid(meta):
         _softsign_glu_bwd_elem_kernel[elem_grid](
             left, gate, grad_y,
             grad_left_pre, grad_gate,
-            M, N, N,
+            M, N,
             left.stride(0), left.stride(1),
             gate.stride(0), gate.stride(1),
             grad_y.stride(0), grad_y.stride(1),

training/variance_task.py

@@ -124,11 +124,13 @@ def __init__(self):
         if hparams.get('use_fused_kernels', False):
             from modules.kernels.integration import patch_variance_model
             from lightning.pytorch.utilities.rank_zero import rank_zero_info
-            n = patch_variance_model(
-                self.model,
-                glu_type=hparams.get('backbone_args', {}).get('glu_type', 'softsign_glu'),
-            )
-            rank_zero_info('Fused kernels: patched %d LYNXNet2 blocks in variance model', n)
+            # Read glu_type from nested predictor configs
+            pitch_glu = hparams.get('pitch_prediction_args', {}).get('backbone_args', {}).get('glu_type', 'softsign_glu')
+            var_glu = hparams.get('variances_prediction_args', {}).get('backbone_args', {}).get('glu_type', 'softsign_glu')
+            assert pitch_glu == var_glu == 'softsign_glu', \
+                f"Fused kernels only support softsign_glu, got pitch={pitch_glu} var={var_glu}"
+            n = patch_variance_model(self.model, glu_type='softsign_glu')
+            rank_zero_info('Fused kernels: patched %d LYNXNet2 blocks in variance model (softsign_glu)', n)
 
 
     def _build_model(self):