feat(moore): support paged flash attention prefill

[voltjia]

Summary

• Rebase the Moore paged-prefill implementation onto the master@879d31cc snapshot.
• Add a native Moore path for flash_attn_varlen_func when the caller supplies a paged KV-cache block_table; dense attention continues to use the shared ATen provider.
• Support FP16/BF16, head dimensions 64/128, MHA/GQA, optional one-dimensional ALiBi, non-contiguous outer strides, empty packed Q, dynamic cumulative lengths, and caller-selected streams.
• Reuse the shared paged-attention prefill kernel without exposing or calling a deprecated operator API.
• Route shared paged-attention output conversion through the existing backend Caster<kDev> abstraction; this PR adds no __CUDA_ARCH__ type-selection branches.
• Split the Torch provider into a common ATen algorithm and backend-specific NVIDIA/Moore policy specializations and translation units selected through DEVICE_LIST.
• Keep the Moore regression cases in the existing tests/test_flash_attn_varlen_func.py module and its existing smoke selector.

Current head: cc42fa99ee4977f80b5cf8af4ff25c47235800ab.

Motivation

InfiniLM's paged FlashAttention prefill path calls the public flash_attn_varlen_func API with a KV-cache block table. The previous Moore provider delegated only dense attention to TorchMusa and rejected this paged form, so Moore inference did not have a complete prefill/decode attention path.

Type of Change

• feat - new provider capability
• test - Moore paged-prefill coverage
• build - shared CUDA/MUSA kernel compatibility
• Breaking change

Platforms Affected

• NVIDIA - shared kernel regression coverage
• Moore Threads - new provider
• PyTorch C++ bindings
• Other backends

Validation

Patch identity after rebase

The feature and test commits are a conflict-free semantic rebase of the hardware-tested prefill commits. git range-diff maps both old commits to d4795d03 and 70e50081 one-for-one. Commit abd8d731 removes the eight architecture-macro type branches and threads the backend device tag to Caster<kDev>. Commit cc42fa99 separates common and backend-specific Torch sources, moves NVIDIA/Moore behavior into backend policy specializations, and consolidates the Moore regression tests into the existing test module.

Focused operator tests

NVIDIA A100-SXM4-80GB (exact head, accelerator-dev/nvidia:latest):
  fresh configure/build/install: PASS
  source graph: NVIDIA flash source included once; Moore flash source absent
  flash_attn_varlen_func: 23 passed, 8 skipped
  paged decode: 12 passed
  paged prefill: 24 passed
  unique total: 59 passed, 8 skipped, 0 failed

MetaX C550 (exact head, PyTorch 2.4.0+metax3.2.1.3):
  fresh smoke configure/build/install: PASS (12/12)
  source graph: NVIDIA flash source absent; Moore flash source absent
  targeted wrapper configure/build/install: PASS (9/9)
  paged decode: 12 passed
  paged prefill: 24 passed
  unique total: 36 passed, 0 skipped, 0 failed

Moore MTT S5000 (exact head):
  fresh core build/install: PASS (15/15)
  fresh Python binding build/install: PASS (3/3)
  source graph: Moore flash source included; NVIDIA flash source absent
  merged flash_attn_varlen_func module: 16 passed, 15 skipped
  Moore stream/current-length + empty-Q regressions: 3 passed
  paged decode: 12 passed
  unique total: 28 passed, 15 skipped, 0 failed

All three platforms exercised the relevant FP16/BF16 paths. NVIDIA skips are four linked-provider cases, three Moore-only regressions, and one two-GPU guard on a single-GPU container. Moore skips are eight unsupported local-window cases, four causal unequal-Q/K cases, and three NVIDIA-only stream/device-guard cases.

The standalone Moore paged-prefill provider is not counted above: its existing CTA8 variant requires 36,896 bytes of shared memory while S5000 permits 28,672 bytes. A no-cache compile at the earlier 70e50081 head produces the same six errors, so this is not introduced by the caster or source-layout cleanup.

InfiniLM integration

The behavior-equivalent pre-layout stack was validated with stacked decode PR #962, InfiniLM #506, InfiniRT 0cdbb169, InfiniCCL's open AllGather/Send/Recv stack, and InfiniCCL #69's Moore architecture fix:

• 13/13 selected real-weight Moore inference commands passed.
• Coverage includes 9g-8B FlashAttention graph batches 1/4/16, 9g-8B paged graph batch 32, Qwen3-32B TP4, Baichuan2 TP2, Llama-3.2-3B, ChatGLM3, InternLM3, GLM-4, MiniCPM4, and MiniCPM4 Eagle speculative decoding.
• Every selected command used segmented graph execution with host_segments > 0; none used whole-engine eager fallback.
• Qwen3-32B TP4 passed communicator initialization and BF16 AllReduce, then completed paged FlashAttention graph generation.

The current exact-head focused tests above revalidate all code paths changed by the layout-only follow-up; the full 13-command InfiniLM matrix was not rerun for that follow-up.

Current CI

At exact head cc42fa99:

• NVIDIA primary/shadow, MetaX primary/shadow, and Ascend primary/shadow pass.
• Both clang-format runs, both Ruff runs, and documentation build pass.
• Moore primary/shadow and Iluvatar primary/shadow remain queued or in progress; the exact-head S5000 results above cover the changed Moore paths.
• Cambricon primary/shadow remain red in the existing generated-wrapper duplicate-__half link area. The base snapshot 879d31cc fails at the same stage (base jobs 100186334980 and 100185959499); this PR does not select either NVIDIA or Moore Torch sources in non-target backend builds.

Notes for Reviewers

• The common aten_impl.h has no kNvidia/kMoore, if constexpr, WITH_*, CUDA_ARCH, or __CUDA_ARCH__ platform branch.
• NVIDIA window mapping lives in src/torch/nvidia/...; TorchMusa window and causal-length restrictions live in src/torch/moore/....
• The Moore slot-8 operator() has the same public signature as its base. It only intercepts calls with a paged block_table; dense calls delegate to the shared ATen base. Its implementation lives in the Moore backend directory.
• Provider declarations remain in the common public header so wrapper generation can discover the active implementation without conditionally including backend headers.
• Source isolation is covered by fresh CMake builds and inspection of the generated unity source graph, rather than a test coupled to the exact spelling of the CMake script.
• The paged path requires causal global attention, head dimension 64 or 128, FP16/BF16, and at most one-dimensional ALiBi. It does not return attention probabilities.
• Cumulative K lengths stay device-resident, so graph replay sees in-place updates without host synchronization or temporary length tensors.
• The shared global-memory kernel is intentional on S5000: the existing D128 8-warp variant requires more static shared memory than the device permits.
• PR feat(moore): support paged flash attention decode #962 is stacked on this head and adds the corresponding paged decode provider.

Performance Impact

N/A. This PR makes a previously unsupported Moore execution path functional and makes no comparative performance claim.