[Feature] Add FFN op for Ascend950 - #15572
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a fused FFN linear operator designed to improve performance on Ascend950 hardware. By fusing the linear operations within the FFN structure, the implementation aims to reduce overhead and improve throughput for transformer-based models. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
github-actions
Bot
added
documentation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [Feature] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the FFN (Feed-Forward Network) operator with an optimized fused implementation for the Ascend950 (arch35) architecture. It includes new kernel tiling strategies, operator definitions, and ACLNN interfaces to support fused GELU, SILU, and SWIGLU activations. Feedback was provided regarding a potential use-after-free vulnerability in the ffn_linear function due to pointer lifetime issues and a request to improve the robustness of attribute access by using names instead of fragile numerical indices.
| std::string actStr(activation.data(), activation.size()); | ||
| auto output_size = ffn_linear_npu_output_size(x_contiguous, w1_contiguous, w2_contiguous, actStr); | ||
| at::Tensor y = at::empty(output_size, x_contiguous.options()); | ||
|
|
||
| const char *activationCStr = actStr.c_str(); | ||
| const bool tokensIndexFlag = false; | ||
|
|
||
| EXEC_NPU_CMD(aclnnFFNV2, | ||
| x_contiguous, w1_contiguous, w2_contiguous, | ||
| expertTokens, | ||
| bias1, bias2, | ||
| scale, offset, | ||
| deqScale1, deqScale2, | ||
| antiquantScale1, antiquantScale2, | ||
| antiquantOffset1, antiquantOffset2, | ||
| activationCStr, inner_precise, tokensIndexFlag, | ||
| y); |
There was a problem hiding this comment.
The activationCStr pointer is derived from a local std::string variable actStr. This pointer becomes invalid once ffn_linear function returns. If EXEC_NPU_CMD involves any asynchronous operations or if the underlying aclnnFFNV2 implementation stores this pointer for later use (e.g., within the returned executor), this will lead to a use-after-free bug. This is a very dangerous pattern.
To ensure safety, the lifetime of the string data must be guaranteed for the entire duration of the asynchronous operation. Please ensure that EXEC_NPU_CMD is fully synchronous and does not hold on to the pointer, or refactor the code to manage the string's lifetime appropriately.
| if (attrsNum <= 2) { | ||
| OP_LOGE(context_->GetNodeName(), "enable_uncache attr not registered, use default."); | ||
| return ge::GRAPH_FAILED; | ||
| } | ||
| // 倒数第2个属性 | ||
| auto enableUncacheAttr = attrs->GetAttrPointer<int64_t>(attrsNum - 2); | ||
| enableUncache_ = (enableUncacheAttr != nullptr && *enableUncacheAttr != 0); | ||
|
|
||
| auto shiftValuePtr = attrs->GetAttrPointer<int64_t>(attrsNum - 1); | ||
| shiftValue_ = shiftValuePtr ? *shiftValuePtr : 0; | ||
| shiftValue_ = shiftValue_ == 0 ? ORI_SHIFT_VALUE : shiftValue_; | ||
| supportMmadS8S4_ = true; |
There was a problem hiding this comment.
Accessing operator attributes by their numerical index from the end of the list (attrsNum - 2, attrsNum - 1) is highly fragile. This approach will break if attributes are added, removed, or reordered in the operator definition, leading to incorrect behavior or crashes.
It is strongly recommended to access attributes by their names to make the code more robust and maintainable. For example, use attrs->GetAttrPointer<int64_t>("enable_uncache") instead of relying on indices.
cfcde10 to
464ecb2
Compare
5380e75 to
c9b48c0
Compare
chenchris2
changed the title
Signed-off-by: chenchris2 <1349418798@qq.com>
c9b48c0 to
ac9e282
Compare
2 participants
What this PR does / why we need it?
Refs #10928
On 950PR/DT there is no official arch35 FFN yet, so the block runs as a small-op chain (linear → act → linear), paying 3 kernel launches plus the activation round-trip through GM. This PR adds a fused FFN operator:
Does this PR introduce any user-facing change?
No — it only adds a new operator; no model path or public API changes.
How was this patch tested?
Accuracy (ATK dual-benchmark: fused op vs small-op chain vs CPU-fp32 golden, level L0): 20/20 model shapes + 200+ generalization cases (incl. non-aligned) PASS.
Performance (kernel-level Task Duration vs small-op chain kernel-time sum): on-par or better on all test cases, up to ~20%; 3 launches → 1.