Skip to content

feat(ascend): enable InfiniOps flash attention - #558

Merged
voltjia merged 1 commit into
refactor/adopt-modern-infini-stackfrom
feat/ascend-flash-attention-minimal-20260903
Sep 9, 2026
Merged

feat(ascend): enable InfiniOps flash attention#558
voltjia merged 1 commit into
refactor/adopt-modern-infini-stackfrom
feat/ascend-flash-attention-minimal-20260903

Conversation

@baominghelly

@baominghelly baominghelly commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Allow the InfiniLM flash-attention backend on Ascend devices.
  • Register the InfiniOps FlashAttnVarlenFunc and FlashAttnWithKvcache adapters for Ascend.
  • Select Ascend's native InfiniOps implementation slot (0) while preserving Iluvatar slot 0, Moore slot 8, and the existing slot 16 providers, including Hygon.
  • Extend the static runtime-contract tests for the Ascend registrations and provider selection.

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 model
  • fix — bug fix
  • perf — performance improvement (no behavioral change)
  • refactor — code restructuring without behavior change
  • test — adding or fixing tests only
  • docs — documentation only
  • build / ci — build system or CI configuration
  • chore — tooling, formatting, or other non-code changes
  • Breaking change

Test Results of Involved Models on Supported Platforms (Please attach screenshots)

Platform Model Test Result
Ascend 910C 9g_8b_thinking Single-request inference smoke test Model loads and executes, but deterministic output is not yet correct: the minimal integration produced - as the first token instead of <think>. See Notes for Reviewers.

Static tests after rebasing onto 64fac2f:

python test/static/test_infinicore_runtime_contracts.py
43 tests passed

python test/static/test_modern_infinicore_compatibility.py
11 tests passed

Build verification previously completed on Ascend 910C before the conflict-only rebase:

infinicore_runtime: passed
_infinicore: passed
_infinilm: passed

Formatting verification:

clang-format 21.1.8: passed
ruff 0.15.20 check/format: passed
GitHub Check Format: passed
GitHub Ruff: passed

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 0 implementation of FlashAttnVarlenFunc requires contiguous K/V and rejects a block_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 FlashAttnWithKvcache registration 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 2013f21 on September 9, 2026.


Checklist

Title, Branch, and Commits

  • PR title follows Conventional Commits.
  • Branch name follows the repository feature-branch convention.
  • Commit message follows Conventional Commits.
  • The branch contains one focused feature commit.
  • The branch is rebased on refactor/adopt-modern-infini-stack at 64fac2f.
  • No fixup, squash, formatting-only, or WIP commits remain.

Scope and Design

  • Changes are limited to Ascend flash-attention dispatch and its static contracts.
  • No debug code or unrelated formatting changes are included.
  • Existing backend dispatch behavior is preserved.

General Code Hygiene

  • git diff --check passes.
  • Comments and error messages are in English.
  • No generated .so files are committed.

C++ and Python Formatting

  • Formatting passes with the formatter versions used by CI.

Testing

  • Static runtime-contract tests passed.
  • The affected C++ and Python extension targets previously built on Ascend 910C.
  • Single-request inference correctness is blocked by the InfiniOps paged-KV provider limitation.
  • Offline performance test was not run pending correctness.
  • Sanity benchmark was not run pending correctness.
  • Service test was not run pending correctness.

Build, CI, and Tooling

  • The affected targets previously built from the isolated validation workspace on Ascend 910C.
  • GitHub CI and Ruff checks passed for commit 2013f21.

Security and Safety

  • No secrets, internal model paths, generated binaries, or customer data are included.
  • No new unsafe pointer operations were introduced.

@baominghelly
baominghelly marked this pull request as ready for review September 4, 2026 02:57
@baominghelly
baominghelly requested a review from voltjia September 4, 2026 02:57
@baominghelly
baominghelly force-pushed the feat/ascend-flash-attention-minimal-20260903 branch 2 times, most recently from e286fea to 3cebd8d Compare September 8, 2026 09:44
@baominghelly
baominghelly force-pushed the feat/ascend-flash-attention-minimal-20260903 branch from 3cebd8d to 2013f21 Compare September 9, 2026 02:28
@voltjia
voltjia merged commit 6966281 into refactor/adopt-modern-infini-stack Sep 9, 2026
6 checks passed
@voltjia
voltjia deleted the feat/ascend-flash-attention-minimal-20260903 branch September 9, 2026 02:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants