Skip to content

Fix FA3 ring attention argument forwarding - #774

Open
zjn20030811 wants to merge 1 commit into
xdit-project:mainfrom
zjn20030811:fix/fa3-ring-scale-forward
Open

Fix FA3 ring attention argument forwarding#774
zjn20030811 wants to merge 1 commit into
xdit-project:mainfrom
zjn20030811:fix/fa3-ring-scale-forward

Conversation

@zjn20030811

Copy link
Copy Markdown

Summary

  • fix the FlashAttention-3 ring wrapper call so optional descaling values are passed through the custom autograd function using positional arguments
  • keep compatibility with yunchang FA3 adapters that do not expose descaling keyword parameters when the values are unset
  • return the correct number of gradients for the expanded autograd input contract

This addresses #616, where xFuserLongContextAttention could fail before entering the FA3 kernel because autograd.Function.apply does not accept keyword arguments.

Validation

  • python -m pytest -q tests/core/test_ring_flash_attn_fa3_args.py (4 passed)
  • python -m pytest -q tests/core/test_ring_flash_attn_fa3_args.py tests/core/test_ring_flash_attn.py (4 passed, 1 skipped because flash-attn is unavailable)
  • python -m compileall -q xfuser/core/long_ctx_attention/ring/ring_flash_attn.py tests/core/test_ring_flash_attn_fa3_args.py
  • git diff --check

Sign up for free to 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