feat(edge_softmax): add edge_softmax and CSRTranspose NPU support for Ascend - #20
Open
xuejiakn wants to merge 4 commits into
Open
feat(edge_softmax): add edge_softmax and CSRTranspose NPU support for Ascend#20xuejiakn wants to merge 4 commits into
xuejiakn wants to merge 4 commits into
Conversation
… Ascend Add Ascend NPU adaptation for edge_softmax operator, enabling GAT/PAGTN graph attention network training on Ascend 910B NPU. New operators: - edge_softmax forward/backward: Ascend C kernel with FullLoad/RowSplit dual-mode, FP32/FP16 support, AR/ARA (num_heads=1/>1) branches - CSRTranspose: Ascend implementation via COOToSR(COOTranspose(CSRToCOO)) chain, reusing existing Ascend kernels Framework integration: - kernel.cc: dispatch edge_softmax by data tensor device (not graph context) to handle NPU tensors with CPU-resident graph topology; add CPU fallback for BackwardSegmentCmp and ScatterAdd (needed by GAT WeightedSumAndMax) - array.cc: add Ascend dispatch branch for CSRTranspose Host code (edge_softmax.cc): - ACL host with forward/backward template specializations (FP32/FP16 x int32/int64) - CSC edge ID remapping via CPU gather/scatter (NPU torch index ops unreliable on DGL blob tensors) - Backward interface adapted to DGL convention (sds = out * grad_out) End-to-end validation: - pytest test_gat.py: 2/3 passed, 1 regression needs 800 epochs (CPU also needs more epochs) - pytest test_pagtn.py: no operator crashes (precision is test-epoch limit) - test_edge_softmax_npu.py: 78 tests, 100% pass (FP32/FP16 x fwd/bwd x 6 graph types x 3 num_heads + edge cases + perf) Co-Authored-By: CANNBot <cannbot@ascend>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add edge_softmax and CSRTranspose NPU support, required by GAT/PAGTN attention normalization.
Changes
Restructure note
Replaces #15[10]. Independent of other PRs (base = master).