feat!: return flash attention KV cache softmax LSE

[voltjia]

Summary

• Keep the basic flash_attn_with_kvcache(q, k_cache, v_cache, out) overload unchanged and out-only.
• Extend only the full scalar/tensor cache_seqlens interfaces with a trailing optional softmax_lse output.
• Require softmax_lse exactly when return_softmax_lse=true, and copy the FlashAttention LSE into the caller-provided buffer.

Motivation

return_softmax_lse was already public, but the corresponding auxiliary output was absent, so the full interface could not represent FlashAttention's conditional return contract. The basic overload should continue to model the default out-only call.

API Alignment

InfiniOps API	Alignment	Reference
flash_attn_with_kvcache(q, k_cache, v_cache, out)	FlashAttention default call returns only out	FlashAttention source
flash_attn_with_kvcache(..., return_softmax_lse, out, softmax_lse?)	FlashAttention returns (out, softmax_lse) only when return_softmax_lse=True	FlashAttention source

Type of Change

• feat
• fix
• perf
• refactor
• test
• docs
• build / ci
• chore
• 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

ssh nvidia, accelerator-dev/nvidia:latest
NVIDIA + Torch targeted integration wheel: passed
Wheel SHA-256: f72af45d2439d0c2459ab08da7461e9635f2c434e22b0ee0cce03fe0c6ff353d
tests/test_generate_wrappers.py + tests/test_generate_torch_ops.py: 35 passed
tests/test_flash_attn_with_kvcache.py -k "not scalar": 7 passed, 5 deselected
clang-format 21.1.8 and Ruff 0.15.22: passed

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	Yes	targeted NVIDIA + Torch wheel build passed	Tensor cache_seqlens, paged, basic/default, and stream paths passed
Iluvatar	No	N/A	N/A
MetaX	No	N/A	N/A
Cambricon	No	N/A	N/A
Moore	No	N/A	N/A
Ascend	No	N/A	N/A

Benchmark / Performance Impact

N/A.

Notes for Reviewers

The basic overload remains out-only. The full overload uses std::optional<Tensor> only for the conditional auxiliary output.

Five scalar-cache_seqlens Python cases were deselected because the current master wrapper generator already classifies that scalar overload as std::optional<Tensor>. A clean master CMake generation reproduced the same binding. Fixing tensor/scalar overload generation is intentionally deferred from this return-contract PR.

The NVIDIA adapter keeps InfiniOps' caller-provided out contract by letting FlashAttention 2.7 produce its internal result and copying it into out; directly passing the external buffer is rejected by FlashAttention in the targeted environment.