feat(ascend): enable InfiniOps flash attention - #558
Open
baominghelly wants to merge 2 commits into
Open
Conversation
baominghelly
marked this pull request as ready for review
September 4, 2026 02:57
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 Moore slot8and the existing slot16providers.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, or Cambricon 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:
Build verification on Ascend 910C:
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 is preserved.CI / ChatOps
GitHub CI run #487 completed successfully for commit
50cd0f4on September 4, 2026.Checklist
Title, Branch, and Commits
refactor/adopt-modern-infini-stackat88e4c6e.Scope and Design
General Code Hygiene
git diff --checkpasses..sofiles are committed.C++ and Python Formatting
scripts/format.pypasses with the same formatter versions used by CI.Testing
Build, CI, and Tooling
50cd0f4.Security and Safety