feat(moore): support flash_attn_varlen_func

[voltjia]

Summary

• Add Moore support for the existing flash_attn_varlen_func public operator through TorchMusa's PrivateUse1 registration of aten::_flash_attention_forward.
• Keep one macro-free ATen implementation shared by NVIDIA and Moore; put only Moore provider capability checks in a file-level Moore adapter.
• Discover and link the installed TorchMusa CMake package so its dispatcher and device-guard registrations are retained; no third-party submodule is added.
• Rebase on feat(nvidia): add flash_attn_with_kvcache #818 while preserving both the NVIDIA FlashAttention 2.7 and Moore TorchMusa dependency paths.

Motivation

The public operator already follows Dao FlashAttention's flash_attn_varlen_func API. TorchMusa exposes the same ATen operator for PrivateUse1, so InfiniOps should rely on ATen dispatch instead of embedding CUDA/MUSA branches in the operator implementation.

N/A - no linked issue.

Torch Backend Architecture

The shared implementation in src/torch/ops/flash_attn_varlen_func is deliberately platform-neutral:

• it contains no WITH_NVIDIA / WITH_MOORE preprocessor branches;
• it includes no CUDA or MUSA headers and creates no vendor stream guards;
• it converts InfiniOps tensors with ToAtenTensor<kDev> and calls at::_flash_attention_forward once;
• it follows the caller's current ATen stream, matching the existing add and gemm Torch backends.

ATen dispatch selects CUDA or PrivateUse1 from the tensor device. Moore's real provider limitations are isolated in src/torch/moore/ops/flash_attn_varlen_func: local windows are rejected, and causal calls require matching per-sequence query/key lengths. The adapter then delegates to the same shared ATen implementation. This is a file-level backend boundary, not an in-file platform branch.

find_package(TorchMusa) and --no-as-needed linkage remain necessary at the CMake layer because TorchMusa registers PrivateUse1 kernels and device guards through library static initialization. WITH_MOORE is also retained only as a build/codegen selection define for Moore explicit instantiations; it is not referenced by the operator implementation.

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

# NVIDIA, head 8fd8f5dd7fca414e1729207f93c3b621351ba749
INFINI_RT_ROOT=/opt/infinirt \
scripts/dev/build.sh nvidia --smoke --jobs 16 --torch-jobs 2 --binding-jobs 2
[build passed]

python -m pytest tests/test_flash_attn_varlen_func.py --devices nvidia -q
28 passed, 2 skipped

python -m pytest tests -m smoke --devices nvidia -q
85 passed, 16 skipped, 12837 deselected

The two focused NVIDIA skips are Moore-only provider capability cases. The focused matrix covers FP16/BF16, head dimensions 64/128, MHA/GQA, global/local causal attention, unequal Q/K lengths, and explicit/implicit current-stream execution.

The final-head Moore smoke-pruned build and focused test also pass in the target TorchMusa environment. The operator's own smoke case is 1 passed, 29 deselected. The broader Moore smoke selection, with the unrelated #818 kvcache module excluded because that target image raises an import-time illegal instruction, is 48 passed, 15 skipped, 2 failed, 5662 deselected; the only failures are pre-existing FP32 GEMM tolerance cases. GitHub Moore main and shadow CI both pass on the same head.

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	Yes	Smoke-pruned build and smoke passed	Focused operator: 28 passed, 2 skipped; no full-suite run
Moore	Yes	Smoke-pruned build, operator smoke, main CI, and shadow CI passed	Target image focused operator: 16 passed, 14 skipped; broader smoke has only two pre-existing FP32 GEMM tolerance failures
Other CI platforms	No	GitHub checks exercise shared wrapper/build integration	No platform-specific implementation was added

Validation environments

NVIDIA host: ssh nvidia
Image: accelerator-dev/nvidia:latest (sha256:dd94fce...54efd)
Hardware: NVIDIA A100-SXM4-80GB
PyTorch: 2.10.0a0+b4e4ee81d3.nv25.12
CUDA: 13.1

Moore host: ssh huangjiacheng@172.22.162.98
Image: registry.mthreads.com/mcconline/inference/vllm:4.3.6_torch2.7.1_20260323_ubuntu
Hardware: 8 x MTT S5000
PyTorch: 2.7.1
TorchMusa: 2.7.1

Benchmark / Performance Impact

N/A - this PR exposes an existing TorchMusa ATen provider and makes no comparative performance claim.

Notes for Reviewers

• Review the common ATen implementation first: there is one at::_flash_attention_forward call and no platform macro or vendor stream code.
• The Moore file is intentionally a thin capability adapter. It prevents known provider semantic mismatches without duplicating the ATen invocation.
• WITH_MOORE=ON plus WITH_TORCH=ON requires an installed TorchMusa package that provides TorchMusaConfig.cmake and registers aten::_flash_attention_forward for PrivateUse1.
• The CMake WITH_MOORE define is required by generated explicit instantiations; removing it produces unresolved generated Moore operator symbols. It does not specialize this operator's behavior.
• No additional FlashAttention smoke case was added: this PR reuses the existing representative flash_attn_varlen_func smoke coverage.