feat: Triton fused SoftSignGLU kernel for LYNXNet2 - #60
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.
- 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
- 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
- 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)
- Remove global TF32 side-effects from acoustic_task / variance_task (allow_tf32 / set_float32_matmul_precision were unrelated to this PR and silently change fp32 training for every user). - Revert 6 config glu_type from softsign_glu back to atanglu. softsign_glu stays available as an explicit opt-in; switching the default is an architectural decision that needs A/B evidence first. - Add use_fused_kernels_variance: false to base.yaml so the key is discoverable instead of buried in a code comment. - Protect kernel import on Windows / no-Triton: try/import triton in fused_linear_softsign_glu.py + capability gate for sm < 70. Guard Fn class + bwd kernel inside if _TRITON_AVAILABLE so the module loads cleanly when Triton is absent. - Handle bias=None in Fn.forward (would AttributeError on .split()). - Remove no-op view in forward(x.dim()>2): left/gate are already [M,N]. - Add one-time-per-(dtype,K) warning when fused falls back to eager (prevents fp32 / unsupported-dtype confusion). - Integration: replace net closure with self.net[] to avoid the deepcopy stale-closure issue (EMA / SWA safety). - Warmup: run no_grad forward only (DDP-safe); drop unused params (glu_type, num_channels) from signature. - Task files: wrap patch calls in try/except ImportError so a missing Triton install logs a clear message and continues in eager mode; log autocast_dtype=None case so fp32 users know fused is inactive. - Test: add assert thresholds so numerical regressions fail loudly.
- VarianceTask: add missing on_fit_start warmup (mirrors AcousticTask, sweeps pitch_predictor and variance_predictor over denoise_fn / velocity_fn). Previously use_fused_kernels_variance patched blocks but never warmed the autotune cache. - _test(): compare fused AND eager against a shared fp32 reference instead of fused-vs-eager fp16. fp16 eager itself deviates ~1e-2 relative from fp32 at K=512 (different accumulation order), so the old absolute threshold conflated rounding with kernel bugs. New criterion: fused error <= 2x eager error. Verified on RTX 2070 (sm_75, torch 2.12.1+cu130, triton 3.2.0): - kernel test: fused error consistently BELOW eager fp16 error (fp32 accumulators), K=256/512/1024 fwd+bwd all pass - integration test: block fwd diff 1.95e-3, grad diff 2.20e-3 - eval-mode fallback matches train-mode fused output - deepcopy safety: copied block uses own weights (stale-closure fixed) - fallbacks: bias=None, CPU, fp32 (one-time warning), bf16-on-Turing - simulated no-Triton import: module loads, eager fallback works - warmup: no_grad forward-only, no grads created; full autocast fp16 training step through patched LYNXNet2 works
If the modules.kernels.integration import fails, rank_zero_info (imported on the following line inside the same try) is also unbound; import it locally in the except clause like variance_task does.
…ory trick Eager (common_layers.py): - SoftSignGLUFunction: ATanGLUFunction-style custom backward. softsign'(x) = (1-|softsign(x)|)^2, so both partials are precomputable in forward; backward is two pure multiplies, saves one tensor vs naive autograd. SoftSignGLU now uses it in training mode. - DoubleSoftSignGLU: y = softsign(out) * softsign(gate). softsign applied to the whole Linear output then split (elementwise => equivalent). Custom Function precomputes both partials into one [.., 2N] buffer. - lynxnet2: register glu_type 'double_softsign_glu'. Kernel (fused_linear_softsign_glu.py): - IS_DOUBLE constexpr on fwd + bwd-elem kernels; Triton compiles two specializations, zero overhead in the single-gate path. - fused_linear_softsign_glu(x, w, b, is_double=False) public API; eager fallback handles both modes. - integration.py: _FUSABLE_GLU_TYPES = (softsign_glu, double_softsign_glu). Verified on RTX 2070: - eager Functions vs naive autograd in fp64: max diff < 2e-15 (both) - double-mode kernel vs fp32 reference: fused error 2-5x SMALLER than eager fp16 (fwd 3.2e-4 vs 1.8e-3) — fp32 accumulators - single-gate regression: unchanged, all pass - block integration both modes: fwd/grad diff ~2e-3, eval fallback exact
One switch for both acoustic and variance training. The original
rationale for a separate variance key ('backbones too small to benefit')
did not hold up in benchmarks: at 24k frames/batch the variance backbone
gains 1.35x (single) / 1.62x (double softsign). The per-predictor
glu_type check already prevents patching non-softsign backbones, so the
extra key only added confusion. base.yaml comment now documents the
actual trade-off (benefit scales with max_batch_frames).
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughLYNXNet2 now supports SoftSignGLU activation and optional Triton-fused Linear+SoftSignGLU execution. Training tasks patch supported backbones and warm up fused kernels on CUDA. Eager execution remains available through device, dtype, and integration fallbacks. ChangesLYNXNet2 fusion
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TrainingTask
participant DiffusionModule
participant FusedLYNXNet2Block
participant fused_linear_softsign_glu
participant Triton
TrainingTask->>DiffusionModule: patch configured backbones
DiffusionModule->>FusedLYNXNet2Block: replace supported blocks
TrainingTask->>FusedLYNXNet2Block: warm up CUDA execution
FusedLYNXNet2Block->>fused_linear_softsign_glu: run fused projections
fused_linear_softsign_glu->>Triton: execute forward and backward kernels
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (10)
modules/kernels/fused_linear_softsign_glu.py (4)
462-466: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd
stacklevel=2to the fallback warning.Ruff reports B028. Without
stacklevel, the warning points at this module instead of the caller, which makes the fallback harder to attribute to a specific model or layer.♻️ Proposed fix
import warnings warnings.warn( f'Fused SoftSignGLU: dtype {x.dtype} not supported for this GPU; ' - f'falling back to eager. (This message is shown once per (dtype, K) pair.)' + f'falling back to eager. (This message is shown once per (dtype, K) pair.)', + stacklevel=2, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modules/kernels/fused_linear_softsign_glu.py` around lines 462 - 466, Update the warnings.warn call in the fused SoftSignGLU fallback path to include stacklevel=2, so the warning points to the caller rather than this module while preserving the existing message and behavior.Source: Linters/SAST tools
396-404: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
torch.splitinstead ofchunkfor consistency with the eager GLU modules.
SwiGLU,ATanGLU, andSoftSignGLUinmodules/commons/common_layers.pyall usetorch.splitwith an explicit size, and the code comments there state thatchunkbreaks ONNX export. This fallback path is training-only today, so behavior is identical. Aligning the split style keeps one convention and protects the function if it is ever reached during export.♻️ Proposed fix
def _eager_linear_softsign_glu(x, weight, bias, is_double=False): """Unfused reference path — used as fallback for unsupported dtypes/GPUs.""" - if bias is not None: - left, gate = F.linear(x, weight, bias).chunk(2, dim=-1) - else: - left, gate = F.linear(x, weight).chunk(2, dim=-1) + y = F.linear(x, weight, bias) + left, gate = torch.split(y, y.size(-1) // 2, dim=-1) if is_double: return F.softsign(left) * F.softsign(gate) return left * F.softsign(gate)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modules/kernels/fused_linear_softsign_glu.py` around lines 396 - 404, Update _eager_linear_softsign_glu to replace both F.linear(...).chunk(2, dim=-1) calls with torch.split using the explicit half-width expected by the eager GLU modules, preserving the existing bias and activation behavior.
496-537: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExtend
_testto cover theis_double=Truepath and the bias gradient.Two gaps in this validation:
- The test only exercises
is_double=False. TheIS_DOUBLE=Truebranches in_fused_linear_softsign_glu_fwd_kernel(Line 183) and_softsign_glu_bwd_elem_kernel(Line 273) implement different math and are never checked against a reference.DoubleSoftSignGLUis a selectableglu_type, so this path can reach production training unverified.- Line 525 resets
x16.gradandw16.gradbut notb16.grad. The fused backward therefore accumulates on top of the eager bias gradient. No assertion readsb16.grad, so the test passes, but the bias gradient is untested and the state is wrong for any future assertion.💚 Proposed fix
- # fp16 fused - x16.grad = w16.grad = None - y_fused = fused_linear_softsign_glu(x16, w16, b16) - y_fused.backward(grad) - fused_fwd = rel_err(y_fused, ref) - fused_dx = rel_err(x16.grad, x32.grad) - fused_dw = rel_err(w16.grad, w32.grad) + eager_db = rel_err(b16.grad, b32.grad) + + # fp16 fused + x16.grad = w16.grad = b16.grad = None + y_fused = fused_linear_softsign_glu(x16, w16, b16) + y_fused.backward(grad) + fused_fwd = rel_err(y_fused, ref) + fused_dx = rel_err(x16.grad, x32.grad) + fused_dw = rel_err(w16.grad, w32.grad) + fused_db = rel_err(b16.grad, b32.grad) assert fused_fwd <= eager_fwd * MARGIN + 1e-6, \ f'K={K} fwd: fused={fused_fwd:.4e} vs eager={eager_fwd:.4e}' assert fused_dx <= eager_dx * MARGIN + 1e-6, \ f'K={K} grad_x: fused={fused_dx:.4e} vs eager={eager_dx:.4e}' assert fused_dw <= eager_dw * MARGIN + 1e-6, \ f'K={K} grad_w: fused={fused_dw:.4e} vs eager={eager_dw:.4e}' + assert fused_db <= eager_db * MARGIN + 1e-6, \ + f'K={K} grad_b: fused={fused_db:.4e} vs eager={eager_db:.4e}'Add an equivalent loop, or an
is_doubleparameter, that builds the reference asF.softsign(l32) * F.softsign(g32)and callsfused_linear_softsign_glu(x16, w16, b16, is_double=True).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modules/kernels/fused_linear_softsign_glu.py` around lines 496 - 537, Extend the validation loop around the existing fp32/eager/fused comparisons to exercise both `is_double=False` and `is_double=True`, using `F.softsign(l32) * F.softsign(g32)` as the double-path reference and passing `is_double=True` to `fused_linear_softsign_glu`. Reset `b16.grad` alongside `x16.grad` and `w16.grad`, and compare the fused bias gradient against `b32.grad` for each mode.
42-47: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeclare Triton as an optional dependency and pin a supported version.
use_fused_kernelsenables Triton kernels, butrequirements.txtdoes not list Triton. Add an optional Triton extra and document the minimum version that supports the used API surface:triton.next_power_of_2,triton.Config(..., num_stages=...), andtl.tensor.Tfor transposes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modules/kernels/fused_linear_softsign_glu.py` around lines 42 - 47, Declare Triton as an optional dependency in the project’s dependency configuration, pinning the minimum supported version required by the fused kernels. Document that this version must support triton.next_power_of_2, triton.Config with num_stages, and tl.tensor.T, while preserving the existing optional-import behavior in the fused kernel module.modules/kernels/integration.py (3)
64-87: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winValidate the
netlayout before you patch the forward pass.
fused_forwardhard-codes indices 0 through 9 ofblock.netand assumes index 4 and index 6 are the two gatednn.Linearlayers.LYNXNet2Block.__init__builds exactly that layout today, but nothing enforces it. If thenn.Sequentialgains, loses, or reorders a layer, this patch either raisesIndexErroror silently computes a different graph while training reports no error.Add a cheap structural check in
wrap_lynxnet2_blockbefore patching, and return the block unpatched if the layout does not match.♻️ Proposed check
is_double = glu_type == 'double_softsign_glu' + # The fused forward hard-codes the Sequential layout built by + # LYNXNet2Block.__init__. Refuse to patch anything else. + net = block.net + if not ( + len(net) == 10 + and isinstance(net[4], nn.Linear) + and isinstance(net[6], nn.Linear) + and net[4].out_features == 2 * net[6].in_features + ): + import warnings + warnings.warn( + 'Unexpected LYNXNet2Block.net layout; leaving block unpatched.', + stacklevel=2, + ) + return block + def fused_forward(self, x):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modules/kernels/integration.py` around lines 64 - 87, Add a structural validation in wrap_lynxnet2_block before replacing fused_forward, confirming block.net has the expected ten-layer layout and that indices 4 and 6 are the gated nn.Linear layers used by fused_forward. If validation fails, return the original block without patching; otherwise preserve the existing wrapping behavior.
280-287: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow the blanket
except Exceptionor log the traceback.Ruff reports BLE001. The handler catches every failure at Line 280, formats only
str(e), and breaks the sweep. A real bug in the fused forward, such as a shape mismatch inwrap_lynxnet2_block, then surfaces as a one-line warning at startup and reappears later as a confusing failure during the first real training step. Log the traceback so the cause is recoverable, and keep the non-fatal behavior.♻️ Proposed change
except Exception as e: # Autotune failure should not crash training — Triton cache # can be built on the first real step instead. - warnings.warn(f'Fused kernel warmup skipped at T={T} ({e})') + import traceback + warnings.warn( + f'Fused kernel warmup skipped at T={T} ({type(e).__name__}: {e})\n' + f'{traceback.format_exc()}', + stacklevel=2, + ) break🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modules/kernels/integration.py` around lines 280 - 287, Update the warmup exception handler in the fused-kernel sweep to preserve non-fatal behavior while recording the full traceback, using traceback-aware warning or logging support instead of only formatting str(e); alternatively narrow the caught exception to the expected autotune/cache failures. Resolve Ruff BLE001 without changing the existing break and cleanup flow around wrap_lynxnet2_block.Source: Linters/SAST tools
174-202: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
patch_acoustic_model/patch_variance_modelhelpers.The training tasks call
patch_diffusion_moduledirectly, not these wrappers.patch_variance_modelalso applies oneglu_typeto both variance predictors and ignores its ownglu_typeargument, so it would not preserve the per-predictor fallback behavior used intraining/variance_task.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modules/kernels/integration.py` around lines 174 - 202, Remove the unused patch_acoustic_model and patch_variance_model helper functions, leaving patch_diffusion_module as the direct integration entry point. Do not alter the existing per-predictor fallback behavior in training/variance_task.py.training/acoustic_task.py (3)
123-141: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe warmup body is duplicated across both task files. The precision parsing, the
autocast_dtypeselection, the fallback log message, and thewarmup_fused_backboneloop are near-identical in the two hooks. The two copies already differ in how they locate the backbones, and any fix to the precision parsing must be applied twice. Move the shared logic into one helper inmodules/kernels/integration.py, for examplewarmup_fused_backbones(backbones, max_frames, precision), and let each task supply only its backbone list.
training/acoustic_task.py#L123-L141: replace the inline body with a call to the shared helper, passing the backbones resolved fromself.model.diffusion.training/variance_task.py#L165-L181: replace the inline body with a call to the same helper, passingself._fused_kernel_backbones.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@training/acoustic_task.py` around lines 123 - 141, The duplicated fused-backbone warmup logic should be centralized in a new helper such as warmup_fused_backbones in modules/kernels/integration.py, including precision parsing, autocast selection, fallback logging, and warmup_fused_backbone iteration. In training/acoustic_task.py lines 123-141, resolve the denoise_fn and velocity_fn backbones and pass them to the helper; in training/variance_task.py lines 165-181, replace the inline logic with the same helper using self._fused_kernel_backbones. Ensure both sites retain their existing backbone-selection behavior.
100-111: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_fused_kernels_fallbackis written but never read.Line 101 and Line 115 maintain this flag, and
training/variance_task.pydoes the same at Line 124 and Line 156. No code reads it. Either use it to gate the warmup and the log message, or remove it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@training/acoustic_task.py` around lines 100 - 111, Remove the unused _fused_kernels_fallback state and its assignments in the acoustic task, and apply the same cleanup to variance_task.py if present. Keep the fused-kernel patching and existing logging behavior unchanged.
123-128: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winRead the effective precision from the trainer, not from
hparams.
hparams['pl_trainer_precision']records the configured value. Insideon_fit_start, the trainer already exists, soself.trainer.precisionreports the precision Lightning resolved at runtime. If the two disagree, the warmup autotunes for the wrong dtype and the warmup no longer protects the first real step.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@training/acoustic_task.py` around lines 123 - 128, Update the precision lookup in on_fit_start to use self.trainer.precision instead of hparams.get('pl_trainer_precision', '32'), while preserving the existing autocast_dtype mapping for 16-bit, bf16, and other precisions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modules/kernels/integration.py`:
- Around line 115-119: Make Triton absence degrade to eager execution
consistently: in modules/kernels/integration.py lines 115-119 within
patch_lynxnet2_model and lines 232-236 within warmup_fused_backbone, replace the
RuntimeError path with a warning and return 0, ensuring the non-CUDA early
return remains reachable. In training/acoustic_task.py lines 112-115 and
training/variance_task.py lines 153-156, retain the existing ImportError
handlers with no direct change because the integration functions will no longer
raise for unavailable Triton.
- Line 338: Remove the unnecessary f-string prefix from the integration success
message print statement, leaving the message text and output behavior unchanged.
- Around line 89-94: Move fused_forward from its local closure into module scope
so it is picklable, and replace the captured is_double value with a block
attribute assigned before monkey-patching forward. Update the bound-method
assignment in the surrounding factory flow to reference the module-level
function while preserving dynamic self binding and fused-kernel behavior.
---
Nitpick comments:
In `@modules/kernels/fused_linear_softsign_glu.py`:
- Around line 462-466: Update the warnings.warn call in the fused SoftSignGLU
fallback path to include stacklevel=2, so the warning points to the caller
rather than this module while preserving the existing message and behavior.
- Around line 396-404: Update _eager_linear_softsign_glu to replace both
F.linear(...).chunk(2, dim=-1) calls with torch.split using the explicit
half-width expected by the eager GLU modules, preserving the existing bias and
activation behavior.
- Around line 496-537: Extend the validation loop around the existing
fp32/eager/fused comparisons to exercise both `is_double=False` and
`is_double=True`, using `F.softsign(l32) * F.softsign(g32)` as the double-path
reference and passing `is_double=True` to `fused_linear_softsign_glu`. Reset
`b16.grad` alongside `x16.grad` and `w16.grad`, and compare the fused bias
gradient against `b32.grad` for each mode.
- Around line 42-47: Declare Triton as an optional dependency in the project’s
dependency configuration, pinning the minimum supported version required by the
fused kernels. Document that this version must support triton.next_power_of_2,
triton.Config with num_stages, and tl.tensor.T, while preserving the existing
optional-import behavior in the fused kernel module.
In `@modules/kernels/integration.py`:
- Around line 64-87: Add a structural validation in wrap_lynxnet2_block before
replacing fused_forward, confirming block.net has the expected ten-layer layout
and that indices 4 and 6 are the gated nn.Linear layers used by fused_forward.
If validation fails, return the original block without patching; otherwise
preserve the existing wrapping behavior.
- Around line 280-287: Update the warmup exception handler in the fused-kernel
sweep to preserve non-fatal behavior while recording the full traceback, using
traceback-aware warning or logging support instead of only formatting str(e);
alternatively narrow the caught exception to the expected autotune/cache
failures. Resolve Ruff BLE001 without changing the existing break and cleanup
flow around wrap_lynxnet2_block.
- Around line 174-202: Remove the unused patch_acoustic_model and
patch_variance_model helper functions, leaving patch_diffusion_module as the
direct integration entry point. Do not alter the existing per-predictor fallback
behavior in training/variance_task.py.
In `@training/acoustic_task.py`:
- Around line 123-141: The duplicated fused-backbone warmup logic should be
centralized in a new helper such as warmup_fused_backbones in
modules/kernels/integration.py, including precision parsing, autocast selection,
fallback logging, and warmup_fused_backbone iteration. In
training/acoustic_task.py lines 123-141, resolve the denoise_fn and velocity_fn
backbones and pass them to the helper; in training/variance_task.py lines
165-181, replace the inline logic with the same helper using
self._fused_kernel_backbones. Ensure both sites retain their existing
backbone-selection behavior.
- Around line 100-111: Remove the unused _fused_kernels_fallback state and its
assignments in the acoustic task, and apply the same cleanup to variance_task.py
if present. Keep the fused-kernel patching and existing logging behavior
unchanged.
- Around line 123-128: Update the precision lookup in on_fit_start to use
self.trainer.precision instead of hparams.get('pl_trainer_precision', '32'),
while preserving the existing autocast_dtype mapping for 16-bit, bf16, and other
precisions.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 9d2b7a82-f5c0-4c9a-95e1-923a1db5d748
📒 Files selected for processing (8)
configs/base.yamlmodules/backbones/lynxnet2.pymodules/commons/common_layers.pymodules/kernels/__init__.pymodules/kernels/fused_linear_softsign_glu.pymodules/kernels/integration.pytraining/acoustic_task.pytraining/variance_task.py
|
@coderabbitai review |
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modules/kernels/fused_linear_softsign_glu.py (1)
364-406: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRestore the one-dimensional input contract.
[K]is valid for the documented[..., K]input shape. The custom function reshapes it to[1, K], but it restores shapes only when the rank is greater than two. Fused CUDA execution then returns[1, N], and backward returns[1, K], instead of the eager[N]output and[K]gradient.Restore the shape for every rank other than two. Add a
[K]forward and backward regression test.Proposed fix
- if x.dim() > 2: + if x.dim() != 2: out = out.view(*orig_shape[:-1], N) ... - if len(ctx.orig_x_shape) > 2: + if len(ctx.orig_x_shape) != 2: grad_x = grad_x.view(*ctx.orig_x_shape)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modules/kernels/fused_linear_softsign_glu.py` around lines 364 - 406, Update FusedLinearSoftSignGLUFn’s forward and backward shape restoration to handle every input rank other than two, including one-dimensional [K] inputs, so fused outputs and gradients match the eager [N] and [K] shapes. Add regression coverage for one-dimensional forward output and backward gradient shapes.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@configs/acoustic.yaml`:
- Around line 6-7: Align the fused-kernel settings in configs/acoustic.yaml
lines 6-7 and configs/variance.yaml lines 6-7: set use_fused_kernels to true and
change each relevant glu_type from atanglu to softsign_glu so both
configurations use the supported fused integration.
---
Outside diff comments:
In `@modules/kernels/fused_linear_softsign_glu.py`:
- Around line 364-406: Update FusedLinearSoftSignGLUFn’s forward and backward
shape restoration to handle every input rank other than two, including
one-dimensional [K] inputs, so fused outputs and gradients match the eager [N]
and [K] shapes. Add regression coverage for one-dimensional forward output and
backward gradient shapes.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 1b104879-003b-4cbb-bad1-09ac3b8fed8e
📒 Files selected for processing (8)
configs/acoustic.yamlconfigs/variance.yamlmodules/backbones/lynxnet2.pymodules/commons/common_layers.pymodules/kernels/fused_linear_softsign_glu.pymodules/kernels/integration.pytraining/acoustic_task.pytraining/variance_task.py
💤 Files with no reviewable changes (1)
- modules/commons/common_layers.py
🚧 Files skipped from review as they are similar to previous changes (2)
- modules/kernels/integration.py
- training/variance_task.py
1 participant
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.
Mirror of openvpi#312, created in this fork for CodeRabbit review.
Summary by CodeRabbit