fix: apply remaining PR openvpi#312 fixes

- weight split at midpoint (support expansion_factor != 1)
- variance nested glu_type from predictor configs
- remove useless warmup from __init__ (model on CPU, no compilation)
- use_fused_kernels default to false

Author: KakaruHayate

configs/base.yaml

@@ -85,7 +85,7 @@ nccl_p2p: true
 ###########
 # fusing
 ###########
-use_fused_kernels: true
+use_fused_kernels: false
 
 ###########
 # finetune

modules/kernels/fused_linear_softsign_glu.py

@@ -39,7 +39,7 @@
 def _fused_linear_softsign_glu_fwd_kernel(
     x_ptr, w_left_ptr, w_right_ptr, b_left_ptr, b_right_ptr,
     y_ptr, left_ptr, gate_ptr,
-    M, K,
+    M, N, K,
     stride_x_b, stride_x_k,
     stride_wl_n, stride_wl_k,
     stride_wr_n, stride_wr_k,
@@ -50,13 +50,15 @@ def _fused_linear_softsign_glu_fwd_kernel(
 ):
     """
     y = (x @ W_left^T + b_left) * softsign(x @ W_right^T + b_right)
-    where softsign(x) = x / (1 + |x|).
 
-    2D grid over (M // BLOCK_M, K // BLOCK_N).
+    N = output dim per GLU half (= inner_dim = dim × expansion_factor)
+    K = input feature dim (= dim for first Linear, inner_dim for second)
+
+    2D grid over (M // BLOCK_M, N // BLOCK_N).
     """
     pid = tl.program_id(0)
     num_pid_m = tl.cdiv(M, BLOCK_M)
-    num_pid_n = tl.cdiv(K, BLOCK_N)
+    num_pid_n = tl.cdiv(N, BLOCK_N)
     pid_m = pid // num_pid_n
     pid_n = pid % num_pid_n
 
@@ -65,9 +67,9 @@ def _fused_linear_softsign_glu_fwd_kernel(
     offs_k = tl.arange(0, BLOCK_K)
 
     m_mask_2d = offs_m[:, None] < M
-    n_mask_nk = offs_n[:, None] < K       # [BLOCK_N, 1] for N×K weight access
-    n_mask_mn = offs_n[None, :] < K       # [1, BLOCK_N] for M×N output access
-    n_mask_1d = offs_n < K
+    n_mask_nk = offs_n[:, None] < N       # [BLOCK_N, 1] for N×K weight access
+    n_mask_mn = offs_n[None, :] < N       # [1, BLOCK_N] for M×N output access
+    n_mask_1d = offs_n < N
 
     acc_left = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
     acc_gate = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
@@ -137,7 +139,7 @@ def _fused_linear_softsign_glu_bwd_kernel(
     left_ptr, gate_ptr, grad_y_ptr,
     grad_x_ptr,
     w_left_ptr, w_right_ptr,
-    M, K,
+    M, K, N,
     stride_l_b, stride_l_n,
     stride_g_b, stride_g_n,
     stride_gy_b, stride_gy_n,
@@ -167,10 +169,10 @@ def _fused_linear_softsign_glu_bwd_kernel(
     k_mask_nk = offs_k[None, :] < K
     acc = tl.zeros([BLOCK_M, BLOCK_K], dtype=tl.float32)
 
-    for n_start in range(0, K, BLOCK_N):
+    for n_start in range(0, N, BLOCK_N):
         n_offs = n_start + offs_n
-        n_mask_mn = n_offs[None, :] < K
-        n_mask_nk = n_offs[:, None] < K
+        n_mask_mn = n_offs[None, :] < N
+        n_mask_nk = n_offs[:, None] < N
 
         left = tl.load(
             left_ptr + offs_m[:, None] * stride_l_b + n_offs[None, :] * stride_l_n,
@@ -279,24 +281,25 @@ class FusedLinearSoftSignGLUFn(torch.autograd.Function):
     @staticmethod
     def forward(ctx, x, weight, bias):
         orig_shape = x.shape
-        K = weight.shape[1]
+        K = weight.shape[1]               # input feature dim (contraction dim)
+        N = weight.shape[0] // 2           # output dim per GLU half
         x_2d = x.reshape(-1, K)
         M = x_2d.shape[0]
 
-        w_left, w_right = weight.split(K, dim=0)
-        b_left, b_right = bias.split(K, dim=0)
+        w_left, w_right = weight.split(N, dim=0)
+        b_left, b_right = bias.split(N, dim=0)
 
-        out = torch.empty(M, K, device=x.device, dtype=x.dtype)
-        left = torch.empty(M, K, device=x.device, dtype=x.dtype)
-        gate = torch.empty(M, K, device=x.device, dtype=x.dtype)
+        out = torch.empty(M, N, device=x.device, dtype=x.dtype)
+        left = torch.empty(M, N, device=x.device, dtype=x.dtype)
+        gate = torch.empty(M, N, device=x.device, dtype=x.dtype)
 
         def grid(meta):
-            return (triton.cdiv(M, meta['BLOCK_M']) * triton.cdiv(K, meta['BLOCK_N']),)
+            return (triton.cdiv(M, meta['BLOCK_M']) * triton.cdiv(N, meta['BLOCK_N']),)
 
         _fused_linear_softsign_glu_fwd_kernel[grid](
             x_2d, w_left, w_right, b_left, b_right,
             out, left, gate,
-            M, K,
+            M, N, K,
             x_2d.stride(0), x_2d.stride(1),
             w_left.stride(0), w_left.stride(1),
             w_right.stride(0), w_right.stride(1),
@@ -306,34 +309,36 @@ def grid(meta):
         )
 
         if x.dim() > 2:
-            out = out.view(*orig_shape[:-1], K)
-            left = left.view(M, K)
-            gate = gate.view(M, K)
+            out = out.view(*orig_shape[:-1], N)
+            left = left.view(M, N)
+            gate = gate.view(M, N)
 
         ctx.save_for_backward(x_2d, weight, left, gate)
         ctx.orig_x_shape = orig_shape
+        ctx.N = N
         return out
 
     @staticmethod
     def backward(ctx, grad_y):
         x, weight, left, gate = ctx.saved_tensors
         M, K = x.shape
-        w_left, w_right = weight.split(K, dim=0)
+        N = ctx.N
+        w_left, w_right = weight.split(N, dim=0)
 
         if grad_y.dim() > 2:
-            grad_y = grad_y.reshape(-1, K)
+            grad_y = grad_y.reshape(-1, N)
 
         # Step 1: Fused element-wise SoftSignGLU backward
-        grad_left_pre = torch.empty(M, K, device=x.device, dtype=x.dtype)
-        grad_gate = torch.empty(M, K, device=x.device, dtype=x.dtype)
+        grad_left_pre = torch.empty(M, N, device=x.device, dtype=x.dtype)
+        grad_gate = torch.empty(M, N, device=x.device, dtype=x.dtype)
 
         def elem_grid(meta):
-            return (triton.cdiv(M, meta['BLOCK_M']) * triton.cdiv(K, meta['BLOCK_K']),)
+            return (triton.cdiv(M, meta['BLOCK_M']) * triton.cdiv(N, meta['BLOCK_K']),)
 
         _softsign_glu_bwd_elem_kernel[elem_grid](
             left, gate, grad_y,
             grad_left_pre, grad_gate,
-            M, K,
+            M, N,
             left.stride(0), left.stride(1),
             gate.stride(0), gate.stride(1),
             grad_y.stride(0), grad_y.stride(1),
@@ -360,7 +365,7 @@ def bwd_grid(meta):
             left, gate, grad_y,
             grad_x,
             w_left, w_right,
-            M, K,
+            M, K, N,
             left.stride(0), left.stride(1),
             gate.stride(0), gate.stride(1),
             grad_y.stride(0), grad_y.stride(1),
@@ -381,20 +386,22 @@ def bwd_grid(meta):
 # ---------------------------------------------------------------------------
 
 def fused_linear_softsign_glu(x, weight, bias):
-    """Fused Linear(2K, K) + SoftSignGLU.
+    """Fused Linear(C, 2*C) + SoftSignGLU, where C = dim (expansion_factor×dim).
 
     y = left * gate / (1 + |gate|)
 
+    Supports expansion_factor != 1 by splitting weight at midpoint.
+
     Args:
-        x: Input [..., K]
-        weight: [2*K, K]
-        bias: [2*K]
+        x: Input [..., K] where K = weight.shape[1] (input dim)
+        weight: [2*N, K] where N = output dim per GLU half (= K × expansion_factor)
+        bias: [2*N]
 
     Returns:
-        [..., K]
+        [..., N]
     """
-    assert weight.shape[0] == 2 * weight.shape[1], \
-        f"Expected [2*K, K], got {weight.shape}"
+    N = weight.shape[0] // 2
+    K = weight.shape[1]
     # Match weight/bias dtype to input (handles 16-mixed precision where
     # weights are fp32 but activations are autocast to fp16)
     if weight.dtype != x.dtype:

training/acoustic_task.py

@@ -144,19 +144,13 @@ def __init__(self):
 
         # ── Fuse LYNXNet2 backbone kernels (in-place) ──
         if hparams.get('use_fused_kernels', False):
-            from modules.kernels.integration import patch_diffusion_module, warmup_fused_backbone
+            from modules.kernels.integration import patch_diffusion_module
             from lightning.pytorch.utilities.rank_zero import rank_zero_info
             n = patch_diffusion_module(
                 self.model.diffusion,
-                glu_type=hparams['backbone_args'].get('glu_type', 'atanglu'),
+                glu_type=hparams['backbone_args'].get('glu_type', 'softsign_glu'),
             )
-            rank_zero_info('Fused kernels: patched %d LYNXNet2 blocks, warming up...', n)
-            if n > 0:
-                backbone = getattr(self.model.diffusion, 'denoise_fn',
-                                   getattr(self.model.diffusion, 'velocity_fn', None))
-                if backbone is not None:
-                    warmup_fused_backbone(backbone)
-                rank_zero_info('Fused kernels: autotune complete')
+            rank_zero_info('Fused kernels: patched %d LYNXNet2 blocks', n)
 
     def _build_model(self):
         return DiffSingerAcoustic(

training/variance_task.py

@@ -132,23 +132,13 @@ def __init__(self):
 
         # ── Fuse LYNXNet2 backbone kernels (in-place) ──
         if hparams.get('use_fused_kernels', False):
-            from modules.kernels.integration import patch_variance_model, warmup_fused_backbone
+            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', 'atanglu'),
-            )
-            rank_zero_info('Fused kernels: patched %d LYNXNet2 blocks in variance model, warming up...', n)
-            if n > 0:
-                glu_type_patched = hparams.get('backbone_args', {}).get('glu_type', 'atanglu')
-                for predictor_attr in ['pitch_predictor', 'variance_predictor']:
-                    predictor = getattr(self.model, predictor_attr, None)
-                    if predictor is None:
-                        continue
-                    backbone = getattr(predictor, 'denoise_fn', None) or getattr(predictor, 'velocity_fn', None)
-                    if backbone is not None:
-                        warmup_fused_backbone(backbone)
-                rank_zero_info('Fused kernels: autotune complete')
+            # 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')
+            n = patch_variance_model(self.model, glu_type=pitch_glu)
+            rank_zero_info('Fused kernels: patched %d LYNXNet2 blocks in variance model', n)
 
     def _build_model(self):
         return DiffSingerVariance(