Skip to content

[Ops] Add triton-ascend l2norm and fused_norm_gate backends for Ascend NPU#1002

Merged
zhiyuan1i merged 4 commits into
fla-org:mainfrom
zheliuyu:main
Jul 3, 2026
Merged

[Ops] Add triton-ascend l2norm and fused_norm_gate backends for Ascend NPU#1002
zhiyuan1i merged 4 commits into
fla-org:mainfrom
zheliuyu:main

Conversation

@zheliuyu

@zheliuyu zheliuyu commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Summary

Add triton-ascend backends for l2norm and fused_norm_gate (LayerNorm/RMSNorm + gate) on Ascend NPU. Wire them into the existing public APIs via @dispatch('modules'), so callers are routed automatically on Ascend without code changes.

  • l2norm: Add fwd/bwd Triton kernels and register l2norm_fwd / l2norm_bwd on TritonAscendBackend
  • fused_norm_gate: Add LayerNorm/RMSNorm + gate kernels (swish/sigmoid), with support for residual, weight/bias, and recompute paths
  • UB management: Integrate ascend_ub_manager ([Fix] fix benchamrk issues and add ub management for npu #1001)—use compute_ub_block_size for dynamic tiling and iter_axis_launch_chunks for grid limits; remove ROW_OFFSET / PROGRAM_OFFSET in favor of host-side tensor slicing

Test plan

  • pytest tests/modules/test_l2norm.py -v (NPU)
  • pytest tests/modules/test_layernorm_gated.py -v (NPU)

@zheliuyu

zheliuyu commented Jul 3, 2026

Copy link
Copy Markdown
Contributor Author

@zhiyuan1i The PR is ready for review — please take a look. Thanks!

@zhiyuan1i zhiyuan1i 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. Verified the l2norm and fused_norm_gate fwd/bwd math (l2norm bwd identity, norm+gate swish/sigmoid derivatives, and the LN/RMSNorm backward reduction), the UB-aware tiling via ascend_ub_manager, and the host-side slice chunking / NS capping — grid axes are all bounded, no unchunked-axis edge here. The @dispatch('modules') wiring and TritonAscendBackend method signatures match the public APIs.

NPU A2 CI passes, and the new backend is exercised by the existing test_l2norm.py / test_layernorm_gated.py via dispatch on NPU, so coverage is good. The red H100 test-models check is the pre-existing tilelang/TVM issue on main (this PR is NPU-only), not a regression.

Non-blocking nit: l2norm_bwd_kernel1 takes an unused eps to match the public signature.

@zhiyuan1i
zhiyuan1i merged commit 6f62104 into fla-org:main Jul 3, 2026
12 of 14 checks passed
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