feat(moore): support paged flash attention decode - #962
Merged
Conversation
voltjia
force-pushed
the
feat/moore-flash-attn-kvcache
branch
from
September 2, 2026 10:14
1b4e3b2 to
6f2a152
Compare
voltjia
force-pushed
the
feat/moore-flash-attn-varlen
branch
from
September 2, 2026 13:34
cd69650 to
70e5008
Compare
voltjia
force-pushed
the
feat/moore-flash-attn-kvcache
branch
from
September 2, 2026 13:34
6f2a152 to
c2d7605
Compare
8 tasks
voltjia
force-pushed
the
feat/moore-flash-attn-varlen
branch
from
September 3, 2026 10:21
6d6a409 to
cc42fa9
Compare
voltjia
force-pushed
the
feat/moore-flash-attn-kvcache
branch
from
September 4, 2026 02:53
c2d7605 to
b11c3cc
Compare
voltjia
force-pushed
the
feat/moore-flash-attn-kvcache
branch
from
September 4, 2026 03:30
b11c3cc to
3128c39
Compare
voltjia
force-pushed
the
feat/moore-flash-attn-kvcache
branch
from
September 4, 2026 04:53
3128c39 to
d6048c2
Compare
voltjia
force-pushed
the
feat/moore-flash-attn-kvcache
branch
from
September 4, 2026 13:30
32647b8 to
65e2293
Compare
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
FlashAttnWithKvcache.tests/test_flash_attn_with_kvcache.py: the NVIDIA oracle is imported lazily, while Moore uses the shared local reference with ALiBi support.Current base:
master@2e8a704703a2b2bae950c2203f7b44ff06d43095(includes merged #819).Current head:
65e2293044121fb2b9f599fae36fb82ca6cea52d.Scope
This provider implements the decode contract used by InfiniLM. It intentionally rejects KV append, rotary embedding, cache remapping/left padding, local windows, softcap, LSE output, and scalar cache lengths. Unsupported forms fail closed before dispatch.
Rebase
After #819 merged, its two prefill commits were dropped and the decode changes were replayed onto current master. The final branch contains three commits:
f249c9fc- native Moore paged decode provider and regular functional testsd6048c21- fail closed for scalar cache lengths65e22930- consolidate Moore coverage into the shared KV-cache test moduleThe former smoke-only commit was removed entirely. The final PR has no diff in
src/CMakeLists.txtortests/conftest.py; relative to master it changes only the Moore provider header/source and the existing shared test module.Validation
Shared tests and no new smoke coverage
tests/test_flash_attn_with_kvcache.py; there is no backend-specific test file.flash_attnis imported only inside CUDA oracle paths, so its absence does not skip the whole shared module on Moore or Cambricon.src/CMakeLists.txtandtests/conftest.pyare identical to master.pytest -m smoke --collect-only tests/test_flash_attn_with_kvcache.pyreports no tests collected (21 deselected).Moore S5000
Validated at exact head
65e2293044121fb2b9f599fae36fb82ca6cea52don an isolated MUSA 4.3.6 environment with InfiniRT6b256e0ab0853d9e2a7ef059bac22c847bd02251:10/10steps, compiledpaged.mu.o, and linked bothlibinfiniops.soand the Python extension.tests/test_flash_attn_with_kvcache.py --devices moore -k paged_moore:9 passed, 12 deselected.9 passed, 12 skipped; non-Moore cases skip individually and the module is not skipped by the NVIDIA oracle dependency.NVIDIA
Validated at exact head
65e2293044121fb2b9f599fae36fb82ca6cea52d:12 passed, 9 skipped.InfiniLM integration
The range-diff-equivalent decode implementation was previously validated with InfiniLM #506 and the Moore runtime/communication prerequisites:
host_segments > 0; no command fell back to wholly eager execution.CI
For exact head
65e2293044121fb2b9f599fae36fb82ca6cea52d:torch_npu 2.9.0.post1, which lacksc10_npu::getStreamFromExternaland does not meet master's2.10.0.post4+requirement. The exact base has the same failure.__halfmultiple-definition link failure. Its shadow task also failed, but the artifact upload ended withECONNRESET, so no inner shadow log was recoverable.Notes for Reviewers
FlashAttnWithKvcacheAPI directly.-std=c++17to Iluvatar CUDA flags #69 for Moore BF16 collectives; that dependency is independent of this operator PR.