Skip to content

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

Open
baominghelly wants to merge 2 commits into
refactor/adopt-modern-infini-stackfrom
feat/ascend-flash-attention-minimal-20260903
Open

feat(ascend): enable InfiniOps flash attention#558
baominghelly wants to merge 2 commits 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 Moore slot 8 and the existing slot 16 providers.
  • 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, or Cambricon 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:

test/static/test_infinicore_runtime_contracts.py: 43 passed
test/static/test_modern_infinicore_compatibility.py: 10 passed

Build verification on Ascend 910C:

infinicore_runtime: passed
_infinicore: passed
_infinilm: passed

Formatting verification:

python scripts/format.py --ref origin/refactor/adopt-modern-infini-stack --check: 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 is preserved.

CI / ChatOps

GitHub CI run #487 completed successfully for commit 50cd0f4 on September 4, 2026.


Checklist

Title, Branch, and Commits

  • PR title follows Conventional Commits.
  • Branch name follows the repository feature-branch convention.
  • Commit messages follow Conventional Commits.
  • The feature and formatting commits are focused and independently reviewable.
  • The branch is rebased on refactor/adopt-modern-infini-stack at 88e4c6e.
  • No fixup, squash, 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

  • scripts/format.py passes with the same formatter versions used by CI.

Testing

  • Static runtime-contract tests passed.
  • The affected C++ and Python extension targets build 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 build from the isolated validation workspace on Ascend 910C.
  • GitHub CI and Ruff checks passed for commit 50cd0f4.

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
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.

1 participant