Skip to content

[NPU]: Add NPU support for the swiglu#995

Merged
Tcc0403 merged 2 commits into
linkedin:mainfrom
jiaqiw09:main
Jan 7, 2026
Merged

[NPU]: Add NPU support for the swiglu#995
Tcc0403 merged 2 commits into
linkedin:mainfrom
jiaqiw09:main

Conversation

@jiaqiw09

@jiaqiw09 jiaqiw09 commented Jan 5, 2026

Copy link
Copy Markdown
Contributor

Summary

Add NPU support for the swiglu.

Details

  • Implements a flattened, grid-stride Triton kernel for SwiGLU forward/backward to improve scalability and reduce launch overhead on Ascend NPUs.
  • Uses UB-aware tiling (compute_default_tiling_strategy) and NPU vector core count to dynamically select block size and grid size for better performance stability.

Testing Done

I tested swiglu by following method and all cases passed:

  • python benchmark/scripts/benchmark_swiglu.py

  • pytest -v test/transformers/test_swiglu.py

  • run make test to ensure correctness

  • run make checkstyle to ensure code style

  • run make test-convergence to ensure convergence

@jiaqiw09

jiaqiw09 commented Jan 5, 2026

Copy link
Copy Markdown
Contributor Author

@Tcc0403 would you mind having a preview?

Comment thread src/liger_kernel/ops/backends/_ascend/ops/swiglu.py Outdated
Comment thread src/liger_kernel/ops/backends/_ascend/ops/swiglu.py Outdated
Comment thread src/liger_kernel/ops/backends/_ascend/ops/swiglu.py Outdated
Comment thread src/liger_kernel/ops/backends/_ascend/ops/swiglu.py Outdated
@jiaqiw09

jiaqiw09 commented Jan 6, 2026

Copy link
Copy Markdown
Contributor Author

@Tcc0403 thx for you review, I have just updated the code.

@Tcc0403 Tcc0403 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM

@jiaqiw09

jiaqiw09 commented Jan 7, 2026

Copy link
Copy Markdown
Contributor Author

@Tcc0403 would you mind having a apporval for workflows? Thx.

@Tcc0403 Tcc0403 merged commit fd8e29b into linkedin:main Jan 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants