mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
16e49d4068
## Summary Adds **LigerCute** — fused expert-parallel MoE kernels for NVIDIA Hopper (`sm_90a`) and Blackwell (`sm_100a`), built on CUTLASS + NVSHMEM. A persistent, warp-specialized kernel overlaps NVSHMEM communication with CUTLASS matmuls: communication warps move remote token tiles while the remaining warps run the expert MLP. Forward and backward are both fused, use statically sized symmetric buffers, and support CUDA Graph capture without a token-capacity limit. The change is almost entirely a **new top-level `liger_cute_kernels/` directory** (88 files) plus a small opt-in Python integration. ### What's added | Path | Contents | |---|---| | `liger_cute_kernels/` | New native CUDA package: `csrc/` (CUTLASS/NVSHMEM MoE kernels), `cmake/`, C++ tests, TVM FFI bindings, its own `setup.py`/`pyproject.toml` | | `src/liger_kernel/ops/cute/` | Opt-in `cute` impl registration + `moe_fused` autograd op | | `src/liger_kernel/transformers/moe.py` | `LigerExpertParallelFusedMoe` `nn.Module` | | `test/transformers/test_moe.py`, `test_cute_moe_autograd.py` | Python-side tests | ### Reported BF16 results Benchmark snapshots from a CUDA 12.9 environment — not performance guarantees for every model or system. | Evaluation | Headline result | |---|---| | Standalone MoE kernels | Up to **32% lower forward latency**, **7% lower backward latency** vs the strongest compared implementation | | Communication-intensive H200 cases | **10–35% higher forward throughput** than Comet; selected backward cases up to **~108% higher** than DeepEP | | Qwen3-30B-A3B on 8×H100 | **2.35× speedup / 57% lower step time** than Megatron; **~17% lower step time** than Transformer Engine | | End-to-end convergence | **5.06% final-loss improvement** over the Megatron baseline | ## Details **This is opt-in and does not affect existing users.** - The top-level `liger_kernel` wheel stays **pure Python/Triton**. It does not build, contain, or link any of this native code. - The native core ships as a **separate `liger_cute_kernels` wheel** (CUDA/torch-version prefixed), built independently. - `liger_kernel/ops/cute/__init__.py` registration is **pure metadata** — it does not import the native extension. The ops module (which does load it) is imported only when the impl is actively selected via `LIGER_KERNEL_IMPL=cute`. - If the lck wheel is absent, everything else in `liger_kernel` keeps working; `is_available()` returns `False` and selection raises a helpful `ImportError`. The core's public ABI is `extern "C"` only (no `std::`/torch types cross it), symbols are hidden except `liger_cute_*`/`__tvm_ffi_*`, and libstdc++/libgcc are linked statically — so the core is ABI-agnostic and not tied to a specific torch wheel or a runtime JIT compile. The Python binding is TVM FFI/DLPack based. Blackwell support keeps each kernel's Hopper (WGMMA) consumer byte-for-byte as `Impl<90>` and adds a UMMA/tcgen05 `Impl<100>` selected by an `int Compute` template parameter; the producer is arch-agnostic and shared. ### Changes outside the new directory Only two, both trivial: - `.gitignore` — 3 lines to ignore `liger_cute_kernels/build*` - `setup.py` — `packages=["liger_kernel"]` → `find_packages(where="src")` See `liger_cute_kernels/README.md` for build instructions and the full architecture writeup. ## Testing Done - C++ correctness tests per kernel under `liger_cute_kernels/tests/cpp/` (`test_mlp1_fused`, `test_mlp2_fused`, `test_mlp2_t_fused`, `test_mlp3`, `test_mlp4`, `test_mlp5_fused`, `test_moe_api_dispatch`), validated against an fp32 CPU reference on identically bf16-rounded inputs. Pass bar: `mean_rel < 1%`, `max_rel < 5%`. - Python binding / build / CUDA-graph / NVSHMEM team-setup tests under `liger_cute_kernels/test/`. - Autograd parity tests in `test/transformers/test_cute_moe_autograd.py` and module tests in `test/transformers/test_moe.py`. - Hardware Type: H100 / H200 / B200 (`sm_90a` + `sm_100a`), CUDA 12.9, CUTLASS 4.4.1 - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: Jari Kolehmainen <[email protected]> Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]> Co-authored-by: Chen <[email protected]> Co-authored-by: Copilot <[email protected]> Co-authored-by: autinn <[email protected]> Co-authored-by: autinn <[email protected]> Co-authored-by: jkolehma <[email protected]>
31 lines
296 B
Plaintext
31 lines
296 B
Plaintext
__pycache__/
|
|
*.egg-info/
|
|
site/
|
|
.cache/
|
|
.venv/
|
|
venv/
|
|
.ipynb_checkpoints/
|
|
.vscode/
|
|
.idea/
|
|
|
|
# Misc
|
|
.DS_Store
|
|
|
|
# Build
|
|
build/
|
|
dist/
|
|
|
|
# Doc Build
|
|
site/
|
|
|
|
# Lockfiles
|
|
uv.lock
|
|
|
|
# Benchmark images
|
|
benchmark/visualizations
|
|
.vscode/
|
|
.coverage
|
|
|
|
# Liger CUTE kernels build artifacts
|
|
liger_cute_kernels/build*
|