Skip to content

feat(moore): support paged flash attention decode - #962

Merged
voltjia merged 3 commits into
masterfrom
feat/moore-flash-attn-kvcache
Sep 4, 2026
Merged

feat(moore): support paged flash attention decode#962
voltjia merged 3 commits into
masterfrom
feat/moore-flash-attn-kvcache

Conversation

@voltjia

@voltjia voltjia commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add a native Moore paged decode provider for FlashAttnWithKvcache.
  • Register it as Moore implementation slot 8.
  • Reuse the internal read-only paged-attention decode kernel without registering or calling the deprecated operator API.
  • Keep backend coverage in the shared tests/test_flash_attn_with_kvcache.py: the NVIDIA oracle is imported lazily, while Moore uses the shared local reference with ALiBi support.
  • Cover FP16/BF16, head dimensions 64/128, MHA/GQA, tensor cache lengths, ALiBi, arbitrary outer strides, dynamic metadata, and external streams.
  • Do not add any smoke build or smoke test coverage.

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 tests
  • d6048c21 - fail closed for scalar cache lengths
  • 65e22930 - consolidate Moore coverage into the shared KV-cache test module

The former smoke-only commit was removed entirely. The final PR has no diff in src/CMakeLists.txt or tests/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

  • Moore-specific cases live in tests/test_flash_attn_with_kvcache.py; there is no backend-specific test file.
  • flash_attn is imported only inside CUDA oracle paths, so its absence does not skip the whole shared module on Moore or Cambricon.
  • The shared local reference now covers ALiBi and is reused by Moore.
  • src/CMakeLists.txt and tests/conftest.py are identical to master.
  • On Moore, pytest -m smoke --collect-only tests/test_flash_attn_with_kvcache.py reports no tests collected (21 deselected).
  • The operator and its tests are not added to the smoke build allowlist or smoke matcher.

Moore S5000

Validated at exact head 65e2293044121fb2b9f599fae36fb82ca6cea52d on an isolated MUSA 4.3.6 environment with InfiniRT 6b256e0ab0853d9e2a7ef059bac22c847bd02251:

  • A fresh non-smoke minimal provider build completed 10/10 steps, compiled paged.mu.o, and linked both libinfiniops.so and the Python extension.
  • tests/test_flash_attn_with_kvcache.py --devices moore -k paged_moore: 9 passed, 12 deselected.
  • Full shared KV-cache module: 9 passed, 12 skipped; non-Moore cases skip individually and the module is not skipped by the NVIDIA oracle dependency.
  • The 9 cases comprise the FP16/BF16 x head-dimension 64/128 x MHA/GQA matrix plus the stream/current-metadata regression.
  • The build and test were pinned to one otherwise-idle physical S5000.

NVIDIA

Validated at exact head 65e2293044121fb2b9f599fae36fb82ca6cea52d:

  • Full shared KV-cache module: 12 passed, 9 skipped.
  • All 12 existing CUDA cases pass; the 9 Moore-only cases skip as intended.
  • Ruff lint, Ruff format, and Python AST checks pass.
  • NVIDIA primary and shadow CI pass.

InfiniLM integration

The range-diff-equivalent decode implementation was previously validated with InfiniLM #506 and the Moore runtime/communication prerequisites:

  • 13/13 selected real-weight Moore inference commands passed.
  • The matrix includes paged prefill and decode, explicit FlashAttention, segmented graph replay, TP2/TP4 BF16 collectives, and MiniCPM4 Eagle speculative decoding.
  • All 13 commands reported graph segments and host_segments > 0; no command fell back to wholly eager execution.

CI

For exact head 65e2293044121fb2b9f599fae36fb82ca6cea52d:

  • Ruff, clang-format, documentation build, NVIDIA primary/shadow, and MetaX primary/shadow pass.
  • Ascend primary/shadow fail during configure because the runner has torch_npu 2.9.0.post1, which lacks c10_npu::getStreamFromExternal and does not meet master's 2.10.0.post4+ requirement. The exact base has the same failure.
  • Cambricon primary reproduces the exact base's generated-instantiation __half multiple-definition link failure. Its shadow task also failed, but the artifact upload ended with ECONNRESET, so no inner shadow log was recoverable.
  • Iluvatar primary/shadow and Moore primary/shadow are still queued on self-hosted CI; exact-head Moore hardware validation is reported above.

Notes for Reviewers

  • This PR does not add a public deprecated attention operator. It implements the current FlashAttnWithKvcache API directly.
  • Scalar cache lengths are rejected deliberately: accepting them would make graph replay observe stale metadata.
  • The InfiniLM integration also needs InfiniCCL fix: add -std=c++17 to Iluvatar CUDA flags #69 for Moore BF16 collectives; that dependency is independent of this operator PR.

@voltjia
voltjia force-pushed the feat/moore-flash-attn-kvcache branch from 1b4e3b2 to 6f2a152 Compare September 2, 2026 10:14
@voltjia
voltjia force-pushed the feat/moore-flash-attn-varlen branch from cd69650 to 70e5008 Compare September 2, 2026 13:34
@voltjia
voltjia force-pushed the feat/moore-flash-attn-kvcache branch from 6f2a152 to c2d7605 Compare September 2, 2026 13:34
@voltjia
voltjia force-pushed the feat/moore-flash-attn-varlen branch from 6d6a409 to cc42fa9 Compare September 3, 2026 10:21
Base automatically changed from feat/moore-flash-attn-varlen to master September 4, 2026 02:18
@voltjia
voltjia requested a review from a team September 4, 2026 02:18
@voltjia
voltjia force-pushed the feat/moore-flash-attn-kvcache branch from c2d7605 to b11c3cc Compare September 4, 2026 02:53
@voltjia
voltjia force-pushed the feat/moore-flash-attn-kvcache branch from b11c3cc to 3128c39 Compare September 4, 2026 03:30
@voltjia
voltjia force-pushed the feat/moore-flash-attn-kvcache branch from 3128c39 to d6048c2 Compare September 4, 2026 04:53
@voltjia
voltjia force-pushed the feat/moore-flash-attn-kvcache branch from 32647b8 to 65e2293 Compare September 4, 2026 13:30
@voltjia
voltjia merged commit 5d32b52 into master Sep 4, 2026
12 of 20 checks passed
@voltjia
voltjia deleted the feat/moore-flash-attn-kvcache branch September 4, 2026 14:46
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