Add support for jamba model with Liger Kernel

[yubofredwang]

Summary

Add support for jamba model with Liger Kernel. The following ops can be patched with Liger kernel:

• RMSNorm
• cross_entropy
• swiglu
• lce_forward

Testing Done

• Hardware Type: A100-80G-PCIe
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[yundai424]

CI is failing, we probably need to set use_mamba_kernels=False in the tests? Or install mamba-ssm in GPU CI

[yubofredwang]

pip install . '[dev]' fails for this PR after mamba-ssm into the dependecies. The reason is that mamba-ssm has a bug in its setup.py that makes it not PEP 517 compliant(basically torch, packaging, wheel has to be installed before pip install mamba-ssm runs, otherwise it complains no module found error). Similar situation applies to causal-conv1d.

There are fixes done on both repo, but never gets merged in or released:

• Make properly pip installable without needing build isolation (uv pip install now works too) state-spaces/mamba#402 (comment)
• Enable this to build with build isolation on Dao-AILab/causal-conv1d#26

My current solution is to comment out the tests until the above issues are fixed. However, I have run the convergence test locally. Any other suggestion is highly welcomed.

[ByronHsu]

this is not good. these two dependencies seem very heavy. is there an alternative?

[yubofredwang]

these won't be installed by default though. only when you pip install . '[test]' they are installed.

they are not heavy actually. installation takes about 15 seconds