Skip to content

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

Open
voltjia wants to merge 3 commits into
feat/moore-flash-attn-varlenfrom
feat/moore-flash-attn-kvcache
Open

feat(moore): support paged flash attention decode#962
voltjia wants to merge 3 commits into
feat/moore-flash-attn-varlenfrom
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.
  • Cover FP16/BF16, head dimensions 64/128, MHA/GQA, tensor cache lengths, ALiBi, arbitrary outer strides, dynamic metadata, and external streams.

Current base: InfiniOps #819 at 70e50081.

Current head: c2d76051ca21774f0b515dfd8afe59537e200ac2.

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.

Stacked Review

This PR contains three decode-only commits on top of #819:

  • ae2ce9c5 - native Moore paged decode provider
  • c6de2f45 - paged decode smoke coverage
  • c2d76051 - fail closed for scalar cache lengths

git range-diff maps the three hardware-tested pre-rebase commits to these current commits one-for-one. After #819 merges, this PR can be retargeted to master without changing its review diff.

Validation

Operator coverage

  • D64/D128 x FP16/BF16 x MHA/GQA.
  • Cache lengths 1/255/256/257, non-sequential block tables, tail -1, ALiBi and scale.
  • Non-contiguous outer strides, cache immutability, empty rows, non-default streams.
  • Same-pointer relaunch after in-place cache-length and block-table updates.
  • Final combined Moore attention smoke: 4 passed, 8 deselected.
  • Moore library and generated Python bindings built successfully.

InfiniLM integration

Validated with InfiniLM #506 and the current 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

Clang-format, Ruff, and documentation checks pass. Because this is a stacked PR whose base is #819 rather than master, the repository backend matrix is not triggered yet; it will run after retargeting.

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