feat: add LigerMLP module

[Pearblossom-M]

Summary

Add LigerMLP, a Triton-based fused SwiGLU MLP that co-optimizes the complete forward and backward dataflow rather than only the element-wise SiLU-and-gating stage.

Linked issue: #1347

Details

When running make test, an xfail occurs. This is not due to an error in the operator implementation, but because in test_misaligned_intermediate_size_not_supported, the intermediate_size is set to 431, which does not satisfy Triton's TensorDescriptor requirement of 16‑byte alignment.
Additionally, since this implementation uses tl.dot, when the input data type is fp32, tl.dot performs computations using tf32 internally, which differs from PyTorch's behavior of using pure fp32 for computation. As a result, the numerical error between this implementation and the reference implementation is relatively larger. Therefore, when the data type is fp32, the tolerance range is set to a larger value.

Testing Done

• Hardware Type: RTX 5060ti 16G（SM120）
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

python -m pytest test/transformers/test_mlp.py
================================================================================================== test session starts ===================================================================================================
platform linux -- Python 3.12.3, pytest-9.0.2, pluggy-1.6.0
rootdir: /home/ma-jh/Liger-Kernel
configfile: pyproject.toml
plugins: rerunfailures-16.4, asyncio-1.4.0, xdist-3.8.0, anyio-4.12.0, cov-7.1.0
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 24 items                                                                                                                                                                                                       

test/transformers/test_mlp.py::test_correctness_llamamlp[dtype0-200.0-0.02-2-256-256-512] PASSED                                                                                                                   [  4%]
test/transformers/test_mlp.py::test_correctness_llamamlp[dtype0-200.0-0.02-6-42-128-432] PASSED                                                                                                                    [  8%]
test/transformers/test_mlp.py::test_correctness_llamamlp[dtype0-200.0-0.02-3-37-96-264] PASSED                                                                                                                     [ 12%]
test/transformers/test_mlp.py::test_correctness_llamamlp[dtype1-10000.0-0.01-2-256-256-512] PASSED                                                                                                                 [ 16%]
test/transformers/test_mlp.py::test_correctness_llamamlp[dtype1-10000.0-0.01-6-42-128-432] PASSED                                                                                                                  [ 20%]
test/transformers/test_mlp.py::test_correctness_llamamlp[dtype1-10000.0-0.01-3-37-96-264] PASSED                                                                                                                   [ 25%]
test/transformers/test_mlp.py::test_correctness_llamamlp[dtype2-100.0-0.01-2-256-256-512] PASSED                                                                                                                   [ 29%]
test/transformers/test_mlp.py::test_correctness_llamamlp[dtype2-100.0-0.01-6-42-128-432] PASSED                                                                                                                    [ 33%]
test/transformers/test_mlp.py::test_correctness_llamamlp[dtype2-100.0-0.01-3-37-96-264] PASSED                                                                                                                     [ 37%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype0-200.0-0.02-2-256-256-512] PASSED                                                                                                             [ 41%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype0-200.0-0.02-6-42-128-432] PASSED                                                                                                              [ 45%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype0-200.0-0.02-3-37-96-264] PASSED                                                                                                               [ 50%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype1-10000.0-0.01-2-256-256-512] PASSED                                                                                                           [ 54%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype1-10000.0-0.01-6-42-128-432] PASSED                                                                                                            [ 58%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype1-10000.0-0.01-3-37-96-264] PASSED                                                                                                             [ 62%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype2-100.0-0.01-2-256-256-512] PASSED                                                                                                             [ 66%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype2-100.0-0.01-6-42-128-432] PASSED                                                                                                              [ 70%]
test/transformers/test_mlp.py::test_correctness_inference_mode[dtype2-100.0-0.01-3-37-96-264] PASSED                                                                                                               [ 75%]
test/transformers/test_mlp.py::test_invalid_hidden_act_raises[gelu] PASSED                                                                                                                                         [ 79%]
test/transformers/test_mlp.py::test_invalid_hidden_act_raises[relu] PASSED                                                                                                                                         [ 83%]
test/transformers/test_mlp.py::test_invalid_hidden_act_raises[tanh] PASSED                                                                                                                                         [ 87%]
test/transformers/test_mlp.py::test_supported_hidden_act[silu] PASSED                                                                                                                                              [ 91%]
test/transformers/test_mlp.py::test_supported_hidden_act[swish] PASSED                                                                                                                                             [ 95%]
test/transformers/test_mlp.py::test_misaligned_intermediate_size_not_supported XFAIL (TMA descriptors require 3 * intermediate_size * elem_bytes to be 16-byte aligned; intermediate_size=431 with bf16 (strid...) [100%]

make checkstyle
ruff check --output-format=concise .; ruff_check_status=$?; \
ruff format --check --diff .; ruff_format_status=$?; \
ruff check . --fix; \
ruff format .; \
if [ $ruff_check_status -ne 0 ] || [ $ruff_format_status -ne 0 ]; then \
        exit 1; \
fi
All checks passed!
354 files already formatted
All checks passed!
354 files left unchanged

[Pearblossom-M]

@Tcc0403 I have completed all the modifications, specifically:

• In test_mlp.py, the cosine similarity judgment method is used, and all tests pass with tolerance = 1e‑5.

• The wrapping approach for the Triton kernel has been modified, pure PyTorch operations are now placed inside the autograd.Function.

• The internal logic of _ensure_tma_compatible has been revised. Previously, the implementation could not catch tensors that did not meet the requirements of TensorDescriptor in a timely manner—it would only raise an error when the TensorDescriptor was being created. Now, it raises an error promptly.

[Tcc0403]

Overall lgtm, you can also check out how liger integrates with hf/transformers via monkey_patch.py as a follow-up PR. cc @vaibhavjindal I just noticed there's a new liger cute working project on MoE, I wonder if there's any upcoming change on integration method along with new kernels.

[Pearblossom-M]

@Tcc0403 Thanks for the review! I'll consider this PR (add LigerMLP module) feature-complete as is, and follow up on the monkey_patch.py integration in a separate PR.

As a quick sanity check for that follow-up, I made a small-scale attempt on the Llama path — swapped LigerSwiGLUMLP for LigerMLP inside apply_liger_kernel_to_llama, and it's passing test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_llama plus the bf16 convergence tests for mini_llama3 (both test_mini_models.py and test_mini_models_with_logits.py).

Before I start that follow-up PR, a couple of scoping questions:

1. Rollout scope: should the follow-up PR start as a Llama-only PoC (validate the approach, then extend to other SwiGLU-family models like Mistral, Qwen2, etc. in later PRs), or would you rather I cover all SwiGLU-family models in one PR from the start?
2. API: should LigerMLP fully replace LigerSwiGLUMLP, or keep LigerSwiGLUMLP as a backward-compatible alias?

Happy to go whichever direction makes sense.

[Tcc0403]

Rollout scope: should the follow-up PR start as a Llama-only PoC (validate the approach, then extend to other SwiGLU-family models like Mistral, Qwen2, etc. in later PRs), or would you rather I cover all SwiGLU-family models in one PR from the start?

Let's cover all at once

API: should LigerMLP fully replace LigerSwiGLUMLP, or keep LigerSwiGLUMLP as a backward-compatible alias?

Keep LigerSwiGLUMLP for bwd-compatibility

[Pearblossom-M]

@Tcc0403 I've just pushed new commits addressing your feedback. Could you please take another look when you have a moment? Thanks!

What this change includes

• Added Falcon H1 support, with a new LigerFalconH1MLP.

• In src/liger_kernel/transformers/monkey_patch.py: a global flag USE_FLASH_SWIGLU controls whether every LigerSwiGLUMLP is swapped for LigerMLP, and LigerFalconH1SwiGLUMLP is swapped for LigerFalconH1MLP. The one exception is apply_liger_kernel_to_llama4: Llama4's Llama4MoE.forward reshapes the hidden states into a 2D [T, hidden] tensor before feeding them to the MLP, and the current fused kernels do not support 2D input, so Llama4 keeps the original LigerSwiGLUMLP.

  I first tried a minimal workaround — unsqueeze(0) to turn the 2D input into 3D, then squeeze(0) on the output — but this crashes at the model level: Llama4 performs an in-place addition out.add_(...), and squeeze(0) returns a view. PyTorch does not allow in-place modification of a view produced by a custom autograd Function. .clone() would fix the view issue, but it adds an extra copy of the full hidden-state tensor, so that is not acceptable either. The plan is therefore to keep Llama4 on the original implementation for now, and later modify the kernel to natively support 2D input.

• Unit tests pass: test/transformers/test_mlp.py and test/transformers/test_monkey_patch.py.

Two points about the make convergence results

• In bf16, everything passes except mini_pixtral (test_mini_models_multimodal), mini_gpt_oss (test_mini_models) and mini_qwen3_moe (test_mini_models_with_logits):

  • The mini_gpt_oss and mini_qwen3_moe failures are not caused by this change — they also fail on a completely unmodified branch.

  • The mini_pixtral failure comes from a rounding-path difference, not a formula error. The reference implementation rounds every intermediate result (gate/up GEMM outputs, SiLU output, etc.) to bf16, while the fused kernel keeps all intermediates in fp32 and rounds only once when writing the final result. Note that this test's "loss" is the sum of the vision encoder's last_hidden_state (not a cross-entropy loss), so values around 1e7 are expected by design — the large magnitude is not a sign of divergence. Measured over 32 steps, the largest relative deviation is 3.18e-2; raising loss_rtol from 1e-2 to 5e-2 makes the test pass with margin (4e-2 is also sufficient, the measured threshold is exactly 3.18e-2).

    Maintainer input is needed on whether this relaxation is acceptable, or whether a different testing approach would be preferred.

• In fp32 there are many failures. The root cause is that tl.dot uses TF32 for fp32 matrix math by default, while the reference implementation calls PyTorch's nn.Linear, which computes in true FP32. This mismatch causes the failures. Triton does provide input_precision="ieee" to force true FP32, but that is not a simple fix: switching to IEEE precision triggers a shared memory overflow (the error reports 114712 bytes needed vs. a hardware limit of 101376 bytes). The reason is that TF32 compiles to Tensor Core MMA instructions with a compact data layout, while IEEE cannot use Tensor Cores and falls back to a software FMA path that needs more shared memory. Even the smallest tile configuration in the current autotune search space overflows, so there is no fallback option. Possible fixes include:

  • Relax the fp32 tolerances
  • Force IEEE precision and re-tune the autotune search space (with a performance cost)
  • Use a different way of testing fp32 correctness

  Which approach to take is up to the maintainers.

The raw failure logs are as follows:

pytest test/convergence/fp32/test_mini_models.py

FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_llama3-32-0.0001-dtype1-1e-08-2e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 9
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_llava-32-0.0001-dtype2-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 33
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_qwen2-32-0.0001-dtype5-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 33
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_qwen2_vl-32-0.0001-dtype9-1e-05-0.1-0.005-1e-05-0.005-1e-05] - AssertionError: [Top k logprobs]Number of mismatched elements: 17046
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_qwen2_5_vl-32-0.0001-dtype10-1e-05-0.1-0.005-1e-05-0.005-1e-05] - AssertionError: [Top k logprobs]Number of mismatched elements: 1894
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_qwen3-32-0.0001-dtype6-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 33
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_mistral-32-0.0001-dtype19-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 17
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_ministral-32-0.0001-dtype20-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 17
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_granite3-32-0.0001-dtype24-1e-08-0.0001-0.04-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 13
FAILED test/convergence/fp32/test_mini_models.py::test_mini_model[mini_exaone4-32-1e-05-dtype34-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 32

10 failed

pytest test/convergence/fp32/test_mini_models_with_logits.py

FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_llama3-32-0.0001-dtype1-1e-08-2e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 10
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_llava-32-0.0001-dtype2-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 32
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_mllama-32-0.0001-dtype3-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Top K Logprobs]Number of mismatched elements: 2
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_qwen2-32-0.0001-dtype5-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 32
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_qwen2_vl-32-0.0001-dtype8-1e-08-2e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 26
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_qwen2_5_vl-32-0.0001-dtype9-1e-08-2e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 24
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_qwen3-32-0.0001-dtype6-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 32
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_mistral-32-0.0001-dtype18-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 18
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_ministral-32-0.0001-dtype19-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 18
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_granite3-32-0.0001-dtype23-1e-08-0.0001-0.005-1e-05-0.005-1e-05] - AssertionError: [Top K Logprobs]Number of mismatched elements: 59
FAILED test/convergence/fp32/test_mini_models_with_logits.py::test_mini_model[mini_exaone4-32-1e-05-dtype32-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 30

11 failed

pytest test/convergence/fp32/test_mini_models_multimodal.py

FAILED test/convergence/fp32/test_mini_models_multimodal.py::test_mini_model_multimodal[mini_qwen2_vl-32-0.0001-dtype0-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 30
FAILED test/convergence/fp32/test_mini_models_multimodal.py::test_mini_model_multimodal[mini_qwen2_5_vl-32-0.0001-dtype5-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 33
FAILED test/convergence/fp32/test_mini_models_multimodal.py::test_mini_model_multimodal[mini_mllama-32-0.0001-dtype8-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 28
FAILED test/convergence/fp32/test_mini_models_multimodal.py::test_mini_model_vision[mini_pixtral-32-0.0001-dtype0-1e-08-1e-05-0.005-1e-05-0.005-1e-05] - AssertionError: [Loss]Number of mismatched elements: 32

4 failed

pytest test/convergence/bf16/test_mini_models.py

FAILED test/convergence/bf16/test_mini_models.py::test_mini_model[mini_gpt_oss-32-1e-05-dtype8-0.05-0.05-0.1-0.1-0.01-0.01] - AssertionError: [Loss]Number of mismatched elements: 1

pytest test/convergence/bf16/test_mini_models_with_logits.py

FAILED test/convergence/bf16/test_mini_models_with_logits.py::test_mini_model[mini_qwen3_moe-32-1e-05-dtype7-0.01-0.2-0.1-0.01-0.01-0.01] - AssertionError: [Loss]Number of mismatched elements: 5

pytest test/convergence/bf16/test_mini_models_multimodal.py

FAILED test/convergence/bf16/test_mini_models_multimodal.py::test_mini_model_vision[mini_pixtral-32-0.0001-dtype0-0.001-0.01-1.0-0.01-0.01-0.01] - AssertionError: [Loss]Number of mismatched elements: 1