fix: pack flash_attn_varlen_func softmax LSE output

[voltjia]

Summary

• Normalize the padded three-dimensional ATen softmax_lse result to the public packed (num_heads, total_q) layout in src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc.
• Normalize the ATen oracle in tests/test_flash_attn_varlen_func.py so the existing unequal-length batch covers both padded and packed ATen layouts.

Motivation

PR #854 validated against a newer PyTorch that returns packed LSE output, but the pinned NVIDIA CI image uses PyTorch 2.5 and returns (batch, num_heads, max_seqlen_q). Copying (2, 4, 5) directly into the public (4, 8) output fails with the reported dimension 2 mismatch.

This fixes the deterministic NVIDIA failure in:

• CI run 30627654008
• CI v2 Shadow run 30627654102

No issue is linked; this PR directly repairs the failing master checks above.

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

Final commit 05987b44, ssh nvidia, nvcr.io/nvidia/pytorch:24.10-py3, PyTorch 2.5.0a0+e000cf0, InfiniRT b548ff5a:

python -m pip install . --no-build-isolation --no-deps \
  --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON \
  --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
Successfully built InfiniOps
Wheel SHA-256: 7df376b9bed9a582f0fce80f86a3702bb492c9603607b6fb435344a2bc3d4fb0

python -m pytest tests/test_flash_attn_varlen_func.py -q
19 passed, 1 skipped in 1.09s

python -m pytest tests -m smoke -n 8 -q
107 passed, 25 skipped in 16.87s

Compatibility build, accelerator-dev/nvidia:latest, PyTorch 2.10.0a0+b4e4ee81d3, InfiniRT b548ff5a:

python -m pip install .[dev] --no-build-isolation \
  --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON \
  --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
Successfully built InfiniOps
Wheel SHA-256: 3c7c683f77c2e43083dfb1d6edbb9f7bf9164713cfa45ccc30f09a0a706d34b7

python -m pytest tests/test_flash_attn_varlen_func.py -q
19 passed, 1 skipped in 1.93s

Formatting:

ruff check tests/test_flash_attn_varlen_func.py
All checks passed!

ruff format --check tests/test_flash_attn_varlen_func.py
1 file already formatted

clang-format --dry-run --Werror src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc
passed

Test Results on Supported Platforms

Platform	Affected	Build / Smoke Result	Full Result / Notes
NVIDIA	Yes	smoke build and smoke tests passed	Targeted test passed on PyTorch 2.5 and 2.10; full suite not run for this scoped compatibility fix
Iluvatar	No	N/A - not affected	N/A - not affected
MetaX	No	N/A - not affected	N/A - not affected
Cambricon	No	N/A - not affected	N/A - not affected
Moore	No	N/A - not affected	N/A - not affected
Ascend	No	N/A - not affected	N/A - not affected

Full `pytest` output (optional)

Not run. The change is isolated to the NVIDIA Torch FlashAttention auxiliary-output compatibility path; the affected-platform smoke suite and the complete target module passed.

Benchmark / Performance Impact

N/A - correctness fix. The PyTorch 2.5 compatibility path performs an on-device mask and packing copy only when auxiliary outputs are requested and ATen returns padded LSE storage. The already-packed path used by newer PyTorch versions remains a direct copy.

Notes for Reviewers

• The public output contract remains (num_heads, total_q); this PR does not change the API or base-class validation.
• Query lengths are derived on-device from cu_seqlens_q, avoiding a host synchronization or per-sequence CPU loop.
• s_dmask remains an empty tensor for the currently required zero-dropout inference path.