[RFC] Full-pipeline SwiGLU MLP with forward/backward co-optimization

[Pearblossom-M]

Summary

I have implemented Flash SwiGLU MLP, a Triton-based SwiGLU MLP that co-optimizes the complete forward and backward dataflow rather than only the element-wise SiLU-and-gating stage.

The current LigerSwiGLUMLP keeps gate_proj, up_proj, and down_proj as nn.Linear modules and applies LigerSiLUMulFunction to the outputs of the first two projections. This implementation explores a broader fusion boundary:

• fuse the gate and up projections with SiLU and element-wise gating;
• store A, G, and U in one contiguous buffer;
• avoid materializing the full dA tensor during backward;
• overwrite expired G/U storage with dG/dU;
• aggregate the dWg and dWu computations into one GEMM;
• compute dI in one fused kernel with a single accumulator;
• use a dedicated no-grad path that does not save G/U.

Across the complete benchmark set, most measured end-to-end training speedups are in the 5%–15% range compared with LigerSwiGLUMLP. The no-grad inference path also uses substantially less intermediate memory, although training-memory behavior is shape-dependent.

• Repository: https://github.com/Pearblossom-M/flash-swiglu-mlp
• Technical report: https://github.com/Pearblossom-M/flash-swiglu-mlp/blob/main/flash-swiglu-mlp-technical-report.en.pdf

Question for maintainers

Would an implementation with this fusion scope be appropriate for Liger Kernel?
I would mainly appreciate feedback on whether this broader fusion scope is interesting and fits the project's goals.

Validation environment

• GPU: RTX 5060 Ti 16 GB (SM120)
• Driver: 580.159.03
• CUDA: 13.1
• PyTorch: 2.13.0
• Triton: 3.7.1
• Liger Kernel: 0.8.1
• Dtypes: BF16 and FP16; the tables below show BF16
• Training shapes:
  • (D, H) ∈ {(1024, 2816), (2048, 5632), (4096, 11008)}
  • B = 8
  • S ∈ {256, 512, 1024, 2048, 4096, 8192}
• Inference shapes:
  • B = 1
  • effective request count S ∈ {1, 2, 4, 8, 16, 32, 64, 128, 256}

The main GEMM kernels use Triton's host-side TensorDescriptor API. I have only benchmarked the implementation on SM120, so I am not yet making portability or performance claims for other GPU architectures.

Training results: forward + backward, BF16, B=8

D → H	S	Flash SwiGLU	Liger	torch.compile	Speedup vs. Liger
1024 → 2816	2048	17.25 ms (49.3 TFLOPS)	19.59 ms	19.68 ms	1.14×
1024 → 2816	8192	70.00 ms (48.6 TFLOPS)	78.93 ms	79.16 ms	1.13×
2048 → 5632	4096	137.96 ms (49.3 TFLOPS)	149.80 ms	149.99 ms	1.09×
4096 → 11008	8192	1049.66 ms (50.7 TFLOPS)	1136.68 ms	1136.79 ms	1.08×

These are representative rows; the linked report contains the complete plots and methodology. Each benchmark uses 10 warm-up iterations and reports the mean of 50 measured iterations.

Inference results

For the smallest configuration (D=1024, H=2816), Flash SwiGLU is slightly slower than Liger and torch.compile. At larger dimensions, the three implementations converge, with Flash SwiGLU marginally ahead in several cases.

A likely contributor at small shapes is the fixed CPU-side overhead of host-side TensorDescriptor construction. CUDA Graphs may reduce that overhead, but I have not validated a CUDA-Graph path, so I do not want to claim a decoding advantage at this stage.

Intermediate-memory results

Inference, BF16, B=1

D → H	S	Flash SwiGLU	Liger	torch.compile
1024 → 2816	256	1.9 MB	4.9 MB	3.5 MB
4096 → 11008	256	7.4 MB	16.1 MB	10.8 MB

Across the tested inference configurations, Flash SwiGLU uses approximately 2.2× less intermediate memory than Liger.

Training, forward + backward, BF16, B=8

D → H	S	Flash SwiGLU	Liger	Difference
4096 → 11008	256	178.0 MB	107.0 MB	+66% for Flash
4096 → 11008	8192	5664.0 MB	6356.0 MB	−11% for Flash

Training memory is shape-dependent. At large H and small S, fixed intermediate/workspace overhead makes this implementation less memory-efficient than Liger. It becomes favorable as S increases. I am including both cases to avoid presenting only the favorable region.

Correctness

I compared against a torch.compile-based reference implementation using torch.allclose:

• FP16: atol=1e-3, rtol=1e-4
• BF16: atol=4e-3, rtol=1e-4

The tests cover the forward output and all four gradients (dI, dWg, dWu, and dWd) at both small and large shapes. All tested comparisons pass.

Distributed-training status

Distributed execution has not been tested yet.

Structurally, the implementation uses a standard torch.autograd.Function and introduces no collectives inside the fused kernels. The arithmetic should remain local to each tensor-parallel shard, with the surrounding framework retaining its existing communication points. However, I do not want to claim DDP, ZeRO, or tensor-parallel support until multi-GPU correctness and performance tests have been completed.

Known limitations and validation gaps

• retain_graph=True and higher-order differentiation through create_graph=True are not supported because backward overwrites saved G/U storage with dG/dU.
• The in-place dGU kernel cannot safely run an independent online autotuning search without restoring its overwritten inputs; it currently reuses the forward kernel's selected configuration.
• Performance has been benchmarked only on one SM120 GPU.
• CUDA-Graph decoding and multi-GPU execution have not been tested.

Thanks for maintaining Liger Kernel. It has been a valuable reference while developing and validating this project, and I would appreciate feedback on whether this design direction fits the repository and is worth exploring further.

[Tcc0403]

Yes, it definitely fit liger's repo. Check #937 for some additional features we might want to add on top of it. Feel free to ping me when your PR is ready for review!

[Pearblossom-M]

@Tcc0403 Thanks for the response!

Before I start the PR, I'd like to check how you want this organized in the repo, since my implementation differs from the current design:

1. My implementation fuses the whole SwiGLU MLP into a single kernel, with an interface quite different from the current LigerSiLUMulFunction. Should I replace the implementation in src/liger_kernel/ops/swiglu.py, or add a new "flash" version alongside it?

2. src/liger_kernel/transformers/swiglu.py has several SwiGLU MLP classes. Should all of these be updated to support the flash version, or should I add a separate set of classes instead? If it's the latter, the upstream callers of these classes would need updating too.

Happy to follow whichever structure you'd prefer — let me know and I'll build the PR around that.

[Tcc0403]

Let's put them in two locations: src/liger_kernel/ops/mlp.py and src/liger_kernel/transformers/mlp.py

Three levels of APIs:

1. op level: a thin triton kernel wrapper functions
2. autograd function: compose your ops in forward and backward methods
3. nn.module: LigerMLP, an interchangable layer which has identical __init__ and foward signatures as hf/transformers xxxMLP, convenient for hf monkeypatch later on.

[Pearblossom-M]

Sounds good, thanks for laying that out — I'll go with ops/mlp.py + transformers/mlp.py, following the three-layer structure (op / autograd.Function / LigerMLP).

One scope question before I start: is it okay if the first PR covers the standard SwiGLU case only (no gate/down multiplier), and Falcon H1 support from #937 gets added as a follow-up once the base implementation lands? Trying to keep the initial review scope smaller.

[Tcc0403]

Sure, we can incrementally support new features

[Pearblossom-M]

@Tcc0403 ready for review whenever you have a chance — let me know if anything needs adjusting. #1357