Skip to content

Add FlashAttention-3 support for CUDA - #293

Closed
werwrewe wants to merge 2 commits into
modelscope:dev-swift-v5from
werwrewe:dev-swift-v5-fa3
Closed

werwrewe wants to merge 2 commits into
modelscope:dev-swift-v5from
werwrewe:dev-swift-v5-fa3

Conversation

@werwrewe

Copy link
Copy Markdown
Contributor

PR type

  • Bug Fix
  • New Feature
  • Document Updates
  • More Models or Datasets Support

PR information

What this PR does

Add FlashAttention-3 support for CUDA via the kernelize config entry, as a pure
leaf replacement:

  • New op flash_attention_3 (kernel/ops/flash_attention3/): registered with
    backends=('cuda',), available only on CUDA with flash-attn-3 installed.

Tests

Qwen3 dense SFT on 2×H800 (SP=2), batch_size=12, sequence length ~4096,
peak memory measured with torch.cuda.max_memory_allocated():

case steady-state single step loss[0] → loss[-1] peak MB
sp2_fa3 1.143 2.2358 → 2.2023 47642
sp2_sdpa_attn 1.255 2.2360 → 2.2031 47642

FA3 gives ~9% per-step speedup over the FA2 baseline with identical peak memory
and matching loss/grad-norm curves.

@werwrewe
werwrewe marked this pull request as draft September 28, 2026 03:00
@werwrewe
werwrewe force-pushed the dev-swift-v5-fa3 branch 2 times, most recently from 6286976 to e1565d5 Compare September 28, 2026 03:12
@werwrewe werwrewe closed this Sep 28, 2026
@werwrewe werwrewe reopened this Sep 28, 2026
@werwrewe
werwrewe marked this pull request as ready for review September 28, 2026 03:42
FA2/FA3/FA4 share transformers' single flash_attention_forward entry and
differ only in the leaf kernel, but AttentionInterface.get_interface is
an exact-key lookup. Registering the SP wrapper only under
flash_attention_2 means an FA3 config silently trains without sequence
parallelism. Register the wrapper under every FlashAttention name and
add CPU-only regression tests.
dadb3f0 moved _dispatch_generation from vllm_sampler_tq.py to
generation_submissions.py but left the test importing it from the old
location, breaking pytest collection on dev-swift-v5.
@werwrewe werwrewe closed this Sep 28, 2026
@werwrewe werwrewe reopened this Sep 28, 2026
@werwrewe
werwrewe marked this pull request as draft September 28, 2026 07:03
@werwrewe werwrewe closed this Sep 28, 2026
@werwrewe
werwrewe deleted the dev-swift-v5-fa3 branch September 28, 2026 07:08
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