[Ops] Add triton-ascend l2norm and fused_norm_gate backends for Ascend NPU#1002
Conversation
|
@zhiyuan1i The PR is ready for review — please take a look. Thanks! |
zhiyuan1i
left a comment
There was a problem hiding this comment.
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.
Summary
Add triton-ascend backends for
l2normandfused_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_fwd/l2norm_bwdonTritonAscendBackendascend_ub_manager([Fix] fix benchamrk issues and add ub management for npu #1001)—usecompute_ub_block_sizefor dynamic tiling anditer_axis_launch_chunksfor grid limits; removeROW_OFFSET/PROGRAM_OFFSETin favor of host-side tensor slicingTest plan
pytest tests/modules/test_l2norm.py -v(NPU)pytest tests/modules/test_layernorm_gated.py -v(NPU)