Liger-Kernel

mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00

Files

T

Chunyu e4831e4d64 [NPU] Add NPU Fused MoE kernel (#1183 )

## Motivation

This pr ports `fused_moe.py` and `fused_moe_kernels.py` to an NPU-affine
implementation while preserving the original math. The computational
definition is unchanged: forward remains `W1 (gate/up) -> SwiGLU -> W2
-> token-weighted gather`, and backward still follows `dA' = dO @ W2^T`
to produce `d_pre_act / dS / dW2 / dX / dW1`.
The main changes are execution-strategy optimizations for NPU.


## Note: Use the Skill

For this fused_moe kernel migration, we followed the skill document from
https://github.com/linkedin/Liger-Kernel/issues/1197.


## Testing Done

- Hardware Type: Ascend 910B2
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence


🤖 Generated with: [cursor](https://cursor.com/).

2026-04-24 05:10:16 +00:00

data

MoE Triton Kernel (#1179 )

2026-04-14 23:49:20 +00:00

scripts

[NPU] Add NPU Fused MoE kernel (#1183 )

2026-04-24 05:10:16 +00:00

__init__.py

first commit

2024-08-06 18:44:05 +00:00

BENCHMARK_GUIDELINES.md

[Benchmark]: Add --sweep-mode and --bt to benchmark CLI. (#1163 )

2026-03-31 07:25:20 +00:00

benchmarks_visualizer.py

[Benchmark]: Add --sweep-mode and --bt to benchmark CLI. (#1163 )

2026-03-31 07:25:20 +00:00

README.md

[Benchmark]: Add --sweep-mode and --bt to benchmark CLI. (#1163 )

2026-03-31 07:25:20 +00:00

README.md

Benchmarking Liger Kernels

Follow these steps to benchmark and visualize kernel performance:

Create a benchmark script
- Add your script under benchmark/scripts/
- Name it according to the kernel (e.g., benchmark_<kernel_name>.py)
Run the benchmark
- Results will be saved to benchmark/data/all_benchmark_data.csv
Example: Benchmarking KTO Loss
```
cd benchmark
python scripts/benchmark_kto_loss.py
```
Visualize results
- Use the visualization script with optional modes:
  - --sweep-mode: Select which sweep data to plot.
    - token_length (default): plots where x-axis is sequence length.
    - model_config: plots where x-axis is model configuration.
  - To target specific operation mode(s), pass --kernel-operation-mode one or more values.
  - If you omit --kernel-operation-mode, the script will:
    - For speed metrics: generate plots for all available modes (forward/backward/full).
    - For memory metrics: generate only the full plot.
Examples:
1. Token-length sweep, specific modes (speed):
```
python benchmarks_visualizer.py \
    --kernel-name kto_loss \
    --metric-name speed \
    --kernel-operation-mode forward backward
```
1. Token-length sweep, all modes (speed):
```
python benchmarks_visualizer.py \
    --kernel-name kto_loss \
    --metric-name speed
```
1. Model-config sweep (speed):
```
python benchmarks_visualizer.py \
    --kernel-name geglu \
    --metric-name speed \
    --sweep-mode model_config
```
1. Memory (always full):
```
python benchmarks_visualizer.py \
    --kernel-name kto_loss \
    --metric-name memory
```
View results
- Generated plots will be saved in benchmark/visualizations/
- Filenames include the sweep mode when specified (e.g. geglu_speed_full_model_config.png)