Skip to content

feat!: return flash_attn_varlen_func auxiliary outputs - #854

Merged
voltjia merged 1 commit into
masterfrom
feat/align-flash-attn-varlen-returns
Jul 31, 2026
Merged

feat!: return flash_attn_varlen_func auxiliary outputs#854
voltjia merged 1 commit into
masterfrom
feat/align-flash-attn-varlen-returns

Conversation

@voltjia

@voltjiavoltjia commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Keep the basic flash_attn_varlen_func(..., out) overload unchanged and out-only.
  • Extend only the full interface with trailing optional softmax_lse and s_dmask outputs.
  • Require both auxiliary outputs exactly when return_attn_probs=true, and populate them in the NVIDIA Torch implementation.

Motivation

return_attn_probs was already public, but the full interface omitted FlashAttention's corresponding auxiliary outputs. The basic overload should continue to model the default out-only call.

API Alignment

InfiniOps APIAlignmentReference
flash_attn_varlen_func(..., out)FlashAttention default call returns only outFlashAttention source
flash_attn_varlen_func(..., return_attn_probs, ..., out, softmax_lse?, s_dmask?)FlashAttention returns (out, softmax_lse, S_dmask) only when return_attn_probs=TrueFlashAttention 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_varlen_func.py: 20 passed
clang-format 21.1.8 and Ruff 0.15.22: passed

Test Results on Supported Platforms

PlatformAffectedBuild / Smoke ResultFull Result / Notes
NVIDIAYestargeted NVIDIA + Torch wheel build passed20 passed
IluvatarNoN/AN/A
MetaXNoN/AN/A
CambriconNoN/AN/A
MooreNoN/AN/A
AscendNoN/AN/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 outputs.

This PR is based directly on master. It contains no Moore files and has no dependency on #819; a future Moore implementation can adapt to the public interface independently.

@voltjia
voltjiaforce-pushed the feat/align-flash-attn-varlen-returns branch from b353df7 to b8055efCompareJuly 29, 2026 09:51
@voltjia
voltjia changed the base branch from feat/moore-flash-attn-varlen to masterJuly 29, 2026 09:55
@voltjia
voltjia requested a review from a teamJuly 29, 2026 09:55
@voltjiavoltjia changed the title feat!: return flash attention varlen auxiliary outputsfeat!: return flash_attn_varlen_func auxiliary outputsJul 31, 2026
@voltjia
voltjia merged commit 3f5d9c4 into masterJul 31, 2026
2 checks passed
@voltjia
voltjia deleted the feat/align-flash-attn-varlen-returns branch July 31, 2026 11:37
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

@voltjia