feat(moore): support paged flash attention prefill

[voltjia]

Summary

• Rebase feat/moore-flash-attn-varlen onto current master (793bc3b9).
• Add a Moore native paged-prefill path for flash_attn_varlen_func when a KV-cache block_table is supplied; keep dense attention on the shared ATen provider.
• Reuse the shared paged-attention prefill kernel for FP16/BF16, head sizes 64/128, GQA/MHA, optional one-dimensional ALiBi, non-contiguous outer strides, and caller-selected streams.
• Read cumulative K lengths directly on device, handle empty packed Q, bound page-table traversal, and support packed-Q grids larger than one grid dimension.
• Preserve NVIDIA behavior while making the shared CUDA kernel headers compile cleanly for Moore.

Motivation

InfiniLM's migrated paged-prefill path calls the public flash_attn_varlen_func API with a KV-cache block table. The previous Moore provider only delegated dense attention to TorchMusa and rejected that paged form, leaving Moore without the attention closure needed by those inference commands.

N/A - no linked issue.

Type of Change

• feat - new feature / new operator / new platform
• fix - bug fix
• perf - performance improvement (no behavioral change)
• refactor - code restructuring without behavior change
• test - adding or fixing tests only
• docs - documentation only
• build / ci - build system or CI configuration
• chore - tooling, formatting, or other non-code changes
• Breaking change

Platforms Affected

• CPU (WITH_CPU)
• NVIDIA (WITH_NVIDIA)
• Iluvatar (WITH_ILUVATAR)
• MetaX (WITH_METAX)
• Cambricon (WITH_CAMBRICON)
• Moore (WITH_MOORE)
• Ascend (WITH_ASCEND)
• PyTorch C++ bindings (WITH_TORCH)
• Build system / CMake / CI
• Python bindings / user-facing API

Smoke Test Result

Validation used InfiniOps head 293e84b2 and InfiniRT master at 6b256e0a.

# Moore MTT S5000
python -m pytest tests/test_flash_attn_varlen_func.py \
  tests/test_flash_attn_varlen_func_moore.py \
  tests/test_paged_attention_prefill_infinilm.py --devices moore -q
16 passed, 39 skipped in 16.21s

python -m pytest tests -m smoke --devices moore -q --ignore=tests/test_gemm.py
91 passed, 34 skipped, 3335 deselected in 21.28s

# NVIDIA A100-SXM4-80GB
python -m pytest tests -m smoke --devices cuda -q
99 passed, 32 skipped, 7924 deselected in 7.10s

python -m pytest tests/test_flash_attn_varlen_func.py \
  tests/test_paged_attention_prefill_infinilm.py --devices cuda -q
23 passed, 29 skipped in 1.15s

The unfiltered Moore smoke selection completed with 93 passed, 34 skipped, 2 failed, 6331 deselected. Both failures are existing FP32 GEMM tolerance cases in tests/test_gemm.py; the attention tests and the remaining smoke set pass.

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	Yes	Build passed; smoke: 99 passed, 32 skipped	Focused attention: 23 passed, 29 skipped
Iluvatar	No	N/A - not affected	No platform-specific implementation changed
MetaX	No	N/A - not affected	No platform-specific implementation changed
Cambricon	No	N/A - not affected	No platform-specific implementation changed
Moore	Yes	Build passed; smoke excluding existing GEMM failures: 91 passed, 34 skipped	Focused attention: 16 passed, 39 skipped; full smoke selection has only two existing FP32 GEMM tolerance failures
Ascend	No	N/A - not affected	No platform-specific implementation changed

Validation environments

NVIDIA: NVIDIA A100-SXM4-80GB
Moore: MTT S5000, MUSA 4.3.1, PyTorch 2.7.1, TorchMusa 2.7.1
InfiniRT: 6b256e0ab0853d9e2a7ef059bac22c847bd02251
NVIDIA wheel SHA256: 074cb410a853adec599810a1c1c7df0a1de7857e3514325ffa047e241cbfc6ea
Moore wheel SHA256: 7f23bdf208e4d05ee5a00965b65b10bb8b239d3e4d9c8d0321a6ab3b590e4246

Benchmark / Performance Impact

N/A - this PR adds functional Moore coverage; it makes no comparative performance claim.

Notes for Reviewers

• Dense Moore attention still uses TorchMusa's ATen provider. The new native path is selected only when block_table is present.
• The paged Moore path currently requires causal global attention, head size 64 or 128, FP16/BF16, and at most one-dimensional ALiBi. It does not return attention probabilities.
• K lengths are represented by a device-side accessor over cu_seqlens_k; no host synchronization or temporary length tensor is introduced, so graph replay observes updated cumulative lengths.
• The D64/D128 Moore launcher intentionally uses the shared global-memory kernel. The existing D128 8-warp variant instantiates 36,896 bytes of static shared memory, above the S5000 limit of 28,672 bytes per block.
• Dedicated Moore tests exercise GQA, ALiBi, FP16/BF16, D64/D128, permuted block tables, padded non-contiguous tensors, a non-default stream, in-place cumulative-length updates with replay, and empty packed Q.
• NVIDIA regression coverage is included because the shared paged-attention kernel headers are touched; NVIDIA's existing scalar BF16 path remains under CUDA-only guards.