Skip to content

[NPU] Add NPU Fused MoE kernel #1183

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zheliuyu wants to merge 5 commits into
base: main
Choose a base branch
from
Open

[NPU] Add NPU Fused MoE kernel #1183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d70e6ef
support npu fused moe kernel
zheliuyu Apr 15, 2026
7ee78a1
optimize fused_moe kernel performance
zheliuyu Apr 20, 2026
dabd601
lint code
zheliuyu Apr 20, 2026
494f027
Merge branch 'main' into main
zheliuyu Apr 21, 2026
5feadf4
fix device.synchronize
zheliuyu Apr 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 14 additions & 3 deletions benchmark/scripts/benchmark_fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
from utils import run_memory_benchmark
from utils import run_speed_benchmark

from liger_kernel.ops.fused_moe import LigerFusedMoEFunction
from liger_kernel.ops import LigerFusedMoEFunction
from liger_kernel.utils import get_total_gpu_memory
from liger_kernel.utils import infer_device

Expand Down Expand Up @@ -157,7 +157,12 @@ def _warmup_liger(T, E, H, intermediate_dim, K, dtype, sweep_dim):
warmup_out = warmup_fn()
warmup_out.sum().backward()
del warmup_out
torch.cuda.synchronize()
if device == "cuda":
torch.cuda.synchronize()
elif device == "npu":
torch.npu.synchronize()
else:
torch.cpu.synchronize()
Comment on lines +160 to +165
Copy link
Copy Markdown
Contributor Author

@zheliuyu zheliuyu Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Tcc0403 Thanks for the suggestion. torch provides a cpu equivalent, so I've added it here.



# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -231,7 +236,13 @@ def _probe():
print(f" warmup E={e_val}...")
_warmup_liger(probe_T, e_val, H, intermediate_dim, K, dtype, sweep_dim="E")

torch.cuda.synchronize()
if device == "cuda":
torch.cuda.synchronize()
elif device == "npu":
torch.npu.synchronize()
else:
torch.cpu.synchronize()

print("Autotune warmup complete.\n")

if args.sweep_dim == "num_tokens":
Expand Down
4 changes: 4 additions & 0 deletions src/liger_kernel/ops/backends/_ascend/ops/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
from liger_kernel.ops.backends._ascend.ops.fused_linear_jsd import LigerFusedLinearJSDFunction
from liger_kernel.ops.backends._ascend.ops.fused_linear_jsd import fused_linear_jsd_backward
from liger_kernel.ops.backends._ascend.ops.fused_linear_jsd import fused_linear_jsd_forward
from liger_kernel.ops.backends._ascend.ops.fused_moe import LigerFusedMoEFunction
from liger_kernel.ops.backends._ascend.ops.fused_moe import compute_routing_metadata
from liger_kernel.ops.backends._ascend.ops.fused_neighborhood_attention import LigerFusedNeighborhoodAttentionFunction
from liger_kernel.ops.backends._ascend.ops.fused_neighborhood_attention import fused_neighborhood_attention_forward
from liger_kernel.ops.backends._ascend.ops.geglu import LigerGELUMulFunction
Expand Down Expand Up @@ -149,6 +151,8 @@
"LigerFusedLinearCrossEntropyFunction",
"fused_linear_cross_entropy_forward",
"fused_linear_cross_entropy_backward",
"LigerFusedMoEFunction",
"compute_routing_metadata",
"LigerMHCCoeffsFunction",
"LigerMHCPreFunction",
"LigerMHCPostResFunction",
Expand Down
Loading
Loading