feat: Triton fused SoftSignGLU kernel for LYNXNet2[WIP]#312
Draft
KakaruHayate wants to merge 10 commits into
Draft
feat: Triton fused SoftSignGLU kernel for LYNXNet2[WIP]#312KakaruHayate wants to merge 10 commits into
KakaruHayate wants to merge 10 commits into
Conversation
Fuse nn.Linear + SoftSignGLU into single Triton kernel launches, reducing HBM traffic by ~5 GB/step on LYNXNet2 acoustic backbone. Supported architecture: LYNXNet2 with SoftSignGLU activation only. Key components: - modules/kernels/fused_linear_softsign_glu.py: Fused forward/backward kernels with weight-split strategy to avoid cross-program communication. SoftSignGLU is numerically exact (no Taylor approximation needed). - modules/kernels/integration.py: Non-invasive monkey-patch with automatic eval-mode fallback for ONNX export. - Autotune warmup at training start to avoid first-step compilation lag. Usage: Set 'use_fused_kernels: true' in config (default: true). ONNX export works automatically via model.eval() fallback.
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d8b261c697
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
- Variance model backbones (K=384/512) are too small for Triton fusion to benefit; skip patching to avoid compilation overhead. - Warmup with M=50000 (matching max_batch_frames) so Triton cache hits on the first real training step instead of recompiling. - Reduce autotune configs from 6+5+4=15 to 3+3+2=8, cutting compile time by ~50%.
Warmup at init time runs before Lightning moves the model to GPU, so Triton cannot compile kernels there. Remove the calls entirely; compilation happens naturally on the first real training step. Also fix default glu_type to softsign_glu in both tasks.
Warmup at init time runs before Lightning moves the model to GPU, so Triton cannot compile kernels there. Remove the calls entirely; compilation happens naturally on the first real training step. Also fix default glu_type to softsign_glu in both tasks.
- P1: keep use_fused_kernels opt-in (false default) for backward compat - P1: split weight at midpoint not at K, supporting expansion_factor != 1 - P2: add N param (GLU output dim) separate from K (input/conraction dim) - P2: save ctx.N for backward wrapper calculations
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9256c0b8dd
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
- 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
|
@autumn-2-net hello? |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9c17e7b28f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
KakaruHayate
force-pushed
the
feat/fused-softsign-glu
branch
from
9c17e7b to
ade184c
Compare
KakaruHayate
added a commit
to KakaruHayate/DiffSinger
that referenced
this pull request
Jul 3, 2026
… fewer autotune configs, realistic warmup M
KakaruHayate
added a commit
to KakaruHayate/DiffSinger
that referenced
this pull request
Jul 3, 2026
- 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
|
@autumn-2-net hello |
- Backward: remove custom Triton GEMM (3.4x slower than cuBLAS), use cuBLAS matmuls with out= preallocation for grad_x/grad_w - Autotune key: bucket M by next_power_of_2 to prevent ~2.6s per-batch re-benchmark under variable frame counts (DsBatchSampler) - Forward configs: add large tiles for Ada/Blackwell (4090/5090) + GROUP_M swizzle for L2 reuse; small tiles retained for Turing - Elem backward kernel: fix key=['M','K'] -> key=['N'], fix param name mismatch (M,K -> M,N at call site) - bf16 fallback: tl.dot bf16 crashes on pre-Ampere (Turing sm_75); detect at runtime, fall back to eager path - Warmup: sweep M buckets under torch.autocast, derive cond hidden size from backbone instead of hardcoding 384 - Integration: restrict to softsign_glu only; glu_type default 'swiglu'; per-predictor patching for variance (no blanket assertion) - Acoustic/variance tasks: on_fit_start warmup hook; glu_type fallback Benchmark (RTX 2070, fp16, K=1024): Before: 0.69x (fwd+bwd), 2669ms on new M (135x slower) After: 1.42x (fwd+bwd), 15ms on new M (1.41x faster)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fuse nn.Linear + SoftSignGLU into single Triton kernel launches, reducing HBM traffic by ~5 GB/step on LYNXNet2 acoustic backbone.
Supported architecture: LYNXNet2 with SoftSignGLU activation only.
Key components:
Usage:
Set 'use_fused_kernels: true' in config (default: true). ONNX export works automatically via model.eval() fallback.