feat(ascend): enable InfiniOps flash attention - #558
Merged
voltjia merged 1 commit intoSep 9, 2026
Merged
Conversation
baominghelly
marked this pull request as ready for review
September 4, 2026 02:57
baominghelly
force-pushed
the
feat/ascend-flash-attention-minimal-20260903
branch
2 times, most recently
from
September 8, 2026 09:44
e286fea to
3cebd8d
Compare
baominghelly
force-pushed
the
feat/ascend-flash-attention-minimal-20260903
branch
from
September 9, 2026 02:28
3cebd8d to
2013f21
Compare
voltjia
approved these changes
Sep 9, 2026
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
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
FlashAttnVarlenFuncandFlashAttnWithKvcacheadapters for Ascend.0) while preserving Iluvatar slot0, Moore slot8, and the existing slot16providers, including Hygon.Motivation
InfiniLM currently rejects the flash-attention backend on Ascend before rank workers are created, even when the required InfiniOps Ascend providers are available. This PR adds the minimal framework wiring needed to dispatch those attention calls to InfiniOps without changing the existing NVIDIA, MetaX, Moore, Cambricon, Iluvatar, or Hygon behavior.
Type of Change
feat— new feature / new modelfix— bug fixperf— performance improvement (no behavioral change)refactor— code restructuring without behavior changetest— adding or fixing tests onlydocs— documentation onlybuild/ci— build system or CI configurationchore— tooling, formatting, or other non-code changesTest Results of Involved Models on Supported Platforms (Please attach screenshots)
9g_8b_thinking-as the first token instead of<think>. See Notes for Reviewers.Static tests after rebasing onto
64fac2f:Build verification previously completed on Ascend 910C before the conflict-only rebase:
Formatting verification:
Offline performance, benchmark sanity, and service tests were not run because model-output correctness is still blocked by the InfiniOps paged-KV limitation described below.
Benchmark / Performance Impact
N/A. This PR only adds platform-gated dispatch wiring.
Notes for Reviewers
This is intentionally a minimal InfiniLM-side integration based on the modern InfiniCore stack.
The current Ascend InfiniOps slot
0implementation ofFlashAttnVarlenFuncrequires contiguous K/V and rejects ablock_table. InfiniLM supplies paged K/V plus a block table during prefill. Release builds can therefore execute without producing a useful assertion, but the output is incorrect. A follow-up InfiniOps change must implement the existing paged-KV contract before this PR is ready to merge.The Ascend
FlashAttnWithKvcacheregistration is included for decode dispatch. Existing behavior for other backends, including the Iluvatar adapter wiring added by #557 and Hygon wiring added by #564, is preserved.CI / ChatOps
GitHub CI run #531 and Ruff run #458 completed successfully for commit
2013f21on September 9, 2026.Checklist
Title, Branch, and Commits
refactor/adopt-modern-infini-stackat64fac2f.Scope and Design
General Code Hygiene
git diff --checkpasses..sofiles are committed.C++ and Python Formatting
Testing
Build, CI, and Tooling
2013f21.Security and Safety