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

## 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.

Start by reviewing the linked flash-swiglu-mlp repository and technical report, then compare the proposed full-pipeline fusion with LigerSwiGLUMLP and the project's goals. There are no repository files or tests named in the issue; done would require a maintainer decision on whether this design direction fits and what validation is needed.

Written by the indexing model from the issue text.