feat(torch): add flash_attn_varlen_func operator

[voltjia]

Summary

• Add the public flash_attn_varlen_func operator with the parameter order and defaults of Dao-AILab FlashAttention FA2.
• Add a handwritten NVIDIA Torch backend that invokes PyTorch's ATen FA2 implementation while preserving InfiniOps device and stream semantics.
• Add focused correctness coverage and include one representative case in the operator-pruned smoke build.

Motivation

This is the first representative FlashAttention operator migrated into InfiniOps. It establishes the public operator boundary without vendoring or exposing a third-party package: callers use the upstream-aligned InfiniOps API, while the first provider is the PyTorch CUDA backend already available to WITH_TORCH builds.

N/A - no linked issue.

API Alignment

The public argument order follows Dao-AILab FlashAttention commit 10846960:

flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q, cu_seqlens_k,
    max_seqlen_q, max_seqlen_k,
    dropout_p=0.0,
    softmax_scale=None,
    causal=False,
    window_size=(-1, -1),
    softcap=0.0,
    alibi_slopes=None,
    deterministic=False,
    return_attn_probs=False,
    block_table=None,
    out,
)

InfiniOps uses its normal trailing out execution argument. The backend maps this API to PyTorch's private aten::_flash_attention_forward operator.

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 (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

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

Remote environment: ssh nvidia, accelerator-dev/nvidia:latest, NVIDIA A100-SXM4-80GB, PyTorch 2.10.0a0+b4e4ee81d3.nv25.12, CUDA 13.1.

# Fresh operator-pruned CPU + NVIDIA + Torch build with
# INFINI_OPS_SMOKE_BUILD=ON, followed by:
python -m pytest tests -m smoke --devices nvidia -q

76 passed, 14 skipped, 12295 deselected

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	Yes	pruned build and smoke passed	focused operator suite: 20 passed
Iluvatar	No	N/A - not affected	N/A
MetaX	No	N/A - not affected	N/A
Cambricon	No	N/A - not affected	N/A
Moore	No	N/A - provider deferred	Current MATE image is incompatible with the pinned MATE target; follow-up PR
Ascend	No	N/A - not affected	N/A

Additional validation

tests/test_generate_wrappers.py
13 passed

tests/test_flash_attn_varlen_func.py --devices nvidia
20 passed

ruff check tests/test_flash_attn_varlen_func.py tests/conftest.py
All checks passed!

ruff format --check tests/test_flash_attn_varlen_func.py tests/conftest.py
2 files already formatted

clang-format 21 --dry-run --Werror <new C++ files>
exit 0

# Installed shared-library checks
readelf -d libinfiniops.so
NEEDED: libtorch_cuda.so
NEEDED: libc10_cuda.so

ctypes.CDLL("libinfiniops.so")  # before importing torch
passed

# GitHub Actions on 56ed35e
Build, Ruff, Clang Format, and all main/shadow platform jobs passed.

Benchmark / Performance Impact

N/A - this PR establishes an API and correctness path; it makes no performance claim. The ATen result is copied into the InfiniOps trailing out tensor.

Notes for Reviewers

• The initial supported subset is inference-oriented: dropout_p=0, softcap=0, no ALiBi slopes, no deterministic/probability return mode, and no paged block_table.
• The implementation is handwritten in implementation slot 8, so scripts/torch_ops.yaml is intentionally unchanged.
• The backend uses a private ATen schema. Its exact 15-argument call shape was checked against PyTorch 2.4, the Moore image's PyTorch 2.5, and the tested NVIDIA image's PyTorch 2.10; this remains an explicit version-coupling trade-off.
• Device guards and external/default CUDA stream behavior have dedicated tests. MHA, GQA, unequal sequence lengths, causal/local masks, FP16/BF16, scale, defaults, and multi-GPU guarding are covered against an independent math reference.
• No Dao FlashAttention source, package, or submodule is vendored. A native provider can be added behind the same public operator later.
• Moore support is intentionally out of scope. accelerator-dev/moore:latest provides MUSA 4.3.1, TorchMUSA 2.5, and MATE 0.0.0; the pinned MATE v0.1.3 target requires newer MUSA/TorchMUSA. It should be added as a separate provider PR after the image/toolchain is upgraded.