Skip to content

[Ascend] Add manual and persistent staged FlashMLA kernels - #13

Closed
Denverjin wants to merge 12 commits into
mainfrom
codex/ascend-flashmla-kernels
Closed

Denverjin wants to merge 12 commits into
mainfrom
codex/ascend-flashmla-kernels

Conversation

@Denverjin

@Denverjin Denverjin commented Oct 21, 2026 •

Copy link
Copy Markdown
Owner

Add Ascend sparse-attention prefill and decode examples with a fully manual pipeline and a persistent T.Stage pipeline. The staged implementation uses compiler-managed versions and synchronization, native pair gathers, high-level QK/PV copies and GEMMs, and FP8/FP4 dequantization.

QK and PV use two Buffer views of one four-version L0B allocation. Their MajorK and MajorMN layouts remain distinct while shared .data provides one version clock and a 64 KiB ring without T.assume_conflict. This requires the Cube-view support in fork PR #11.

This PR includes these compiler changes:

Remaining compiler prerequisites:

The placement prerequisite now supports UB, L1, L0A/L0B/L0C, and SS through T.annotate_buffer_offsets, with an explicit fixed version count required for each placed buffer. The staged kernel still uses the earlier T.annotate_l1_offsets API; integrating the new prerequisite also requires migrating that call and adding an explicit fixed q_l1 version count on the prefetched-query path.

The included performance records identify the historical integrated compiler. They are not measurements of this partially integrated branch.

@Denverjin Denverjin closed this Oct 22, 2026
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