Update stable FA3 API (flash_api_stable.cpp) to be inline with the unstable ABI port (flash_api.cpp) by cleonard530 · Pull Request #152 · vllm-project/flash-attention

cleonard530 · 2026-06-18T19:48:59Z

Updated flash_api_stable.cpp to be inline with the current flash_api.cpp and flash_api_torch_lib.cpp.

To view the diff between flash_api_stable.cpp and flash_api.cpp, checkout commit 7da1d8c, where the flash_api_stable.cpp was copied to flash_api.cpp in order to see the diff (i.e. cp flash_api_stable.cpp flash_api.cpp)

Note: The changes in flash_api_stable.cpp are to align it with flash_api.cpp so looking at the commit above is probably more useful than looking at the Files Changed.

A follow-up PR to update CMakeLists.txt to use the flash_api_stable instead of flash_api is planned.

cc @Harry-Chen @janeyx99

vLLM Test Plan/Result

(built using flash_api_stable.cpp instead of flash_api.cpp and flash_api_torch_lib.cpp)

pytest tests/kernels/attention/test_flash_attn.py -k "3"

Test Result

The same test are passing on this branch that are passing on main (and the same failures too)

Migration progress of vLLM using the Audit Python extension torch-abi-audit:

-- extensions --
    [STABLE  ] [abi3-ok               ] _C_stable_libtorch.abi3.so  (stable_shim=87, unstable=0)
    [UNSTABLE] [abi3-ok               ] _flashmla_C.abi3.so  (stable_shim=0, unstable=72)
    [UNSTABLE] [abi3-ok               ] _flashmla_extension_C.abi3.so  (stable_shim=0, unstable=68)
    [STABLE  ] [abi3-ok               ] _moe_C_stable_libtorch.abi3.so  (stable_shim=72, unstable=0)
    [NO-TORCH] [abi3-ok               ] cumem_allocator.abi3.so
    [NO-TORCH] [abi3-ok               ] spinloop.abi3.so
    [UNSTABLE] [uses-private-api      ] third_party/deep_gemm/_C.cpython-312-x86_64-linux-gnu.so  (stable_shim=0, unstable=57)
    [UNSTABLE] [abi3-ok               ] vllm_flash_attn/_vllm_fa2_C.abi3.so  (stable_shim=0, unstable=84)
    [STABLE  ] [abi3-ok               ] vllm_flash_attn/_vllm_fa3_C.abi3.so  (stable_shim=65, unstable=0)

cleonard530 · 2026-06-22T19:38:06Z

-        std::optional<int> page_size,
-        int max_seqlen_k_new,  // 0 means we're not appending new KV
+        int64_t batch_size,
+        int64_t max_seqlen_q,


StableIValues API does not support to<int> but does support to<int64_t> so we update the signatures here as well as the boxed kernels boxed_mha_... below. They get narrowed back to int when assigned to params.<var> in this function below.

cleonard530 · 2026-06-22T19:53:49Z

+        torch::headeronly::ScalarType qkv_dtype,
+        Tensor seqused_k, // b
+        std::optional<Tensor> cu_seqlens_q_,  // b+1
+        std::optional<Tensor> cu_seqlens_k_,  // b+1


Similarly, StableIValue doesn't have specialization for std::optional<const Tensor&> so it is changed here to std::optional. Tensor data is never copied though as Tensor is essentially just a pointer with a reference count.

cleonard530 · 2026-06-22T20:18:15Z


-    auto opts = q.options();


No longer needed since we use new_empty below.

cleonard530 · 2026-06-23T13:12:58Z

        // if (out_.has_value()) { out_.value().copy_(out); }
    }

    return {out, softmax_lse};
 }

-#ifndef FLASHATTENTION_DISABLE_PYBIND
+void boxed_mha_fwd(


Everything below here mirrors the the registration in hopper/flash_api_torch_lib.cpp

cleonard530 · 2026-06-23T20:42:20Z

                      float softmax_scale,
                      int window_size_left,
                      int window_size_right,
-                      int attention_chunk,


attention_chunk was removed in flash_api.cpp

cleonard530 · 2026-06-23T20:42:59Z

    params.deterministic = deterministic;
 }

-template <int Arch, int Split, bool PagedKVNonTMA, bool PackGQA, bool Has_softcap>


Major refactoring here to align with flash_api.cpp

cleonard530 · 2026-06-23T20:58:04Z

+        "Tensor q,"
+        "Tensor k,"
+        "Tensor v,"
+        "Tensor(k_new!)? k_new = None,"


The default parameters here have been removed to align better with legacy code (see total file changes) and in particular, so that softmax_scale would be a required parameter like it is in flash_api.cpp.

…_stable.cpp. Updated flash_api_stable.cpp to be inline with the current flash_api.cpp and flash_api_torch_lib.cpp. And updated static_switch to use a stdexcept instead of a c10/util/Exception.h because importing that header was casuing one unstable symbol (from TORCH_CHECK). Signed-off-by: Chris Leonard <chleonar@redhat.com>

…_api.cpp. Returned the CMakeLists.txt and static_switch.h files back to main, will migrate them in a different PR. Signed-off-by: Chris Leonard <chleonar@redhat.com>

… to see in github. Will move back in next commit Signed-off-by: Chris Leonard <chleonar@redhat.com>

Signed-off-by: Chris Leonard <chleonar@redhat.com>

…iff is easy to see in github. Will move back in next commit Signed-off-by: Chris Leonard <chleonar@redhat.com>

Signed-off-by: Chris Leonard <chleonar@redhat.com>

Harry-Chen · 2026-06-24T02:09:13Z

IIUC, are you actually porting our changes to FA3 (in our fork) to the stable API implementation (from upstream code)?

* Fused Bwd (vllm-project#137) * Fused with Good perf and stride fixed Fix fused bugs isolate failing case fix bug bring back test cases rm split impl in fused use exp2 is global variable now try oom fix save make fused the default limit to reproduce failure return default to split fix head size bug use exp2 back to true * new grid * BLK_SLICE_FACTOR = 1 * add tflops * new commit * test in parrallel * strides added by jusson * disable alibi * fix bugs again * default to fused * add bwd options for varlen * backend filter * default to jingning and batch 4 * best fwd config * fix TRITON_PRINT_AUTOTUNING flag bug * tune * Tuning fwd prefill * add if else * use flag * Minor mask fix * FLIP GRID * use best config for default * print when autotuning * test bfloat16 * fix k and v stride bugs * skip bfloat16 * test kvpacked * disable internal tests * pick default config based on arch * Add alibi in the new bwd kernel (vllm-project#139) * enable alibi for jinging kernel enable alibi for jinging kernel match * save bad configs * fix alibi and causal bug * disable autotune by default * auto tune when benching is good * set best config * remove env var * Update amd_tests.yml * upgrad to triton==3.3.0 * increase shm * use 64 x 64 for now * save * handle 1d alibi * Add fp8 to fused kernel (vllm-project#140) * fp8 stuff find test case compute delta fp8 basic fp8 config passing non causal path works * isolate bad case * fix fp8 bug * didnot fix fp8 bug * back to failing test * fp8 tests passing * skip * skip ref tests --------- Co-authored-by: Aliasger Zaidy <aliasger.zaidy@amd.com> * head, seq, batch (vllm-project#141) * Fix keys (vllm-project#144) * save * rm keys * fix keys * use GHA_RENDER_DEVICES * normal docker * Pad LSE (vllm-project#148) * add round multiple * fix fwd * backward fix * use rounded lse flag * passing ROUNDED_LSE * default is new rounded mode * rename to fused_atmoics and fused_no_atomics * add test for torch_compile * add varlen torch compile test * add old one kernel for ref * fix varlen mismatch bug * fix shape issue in varlen but mismatch * sync torch compile kernel launch * simple varlen test * add debug code * rm old * ignore old impls * DEBUG flag works in interface only * ref uses the righ shape for lse * rm oldest bwd kernel * fix typo * fix varlen bug * fix bug. Get info from q for now * simple shape and stride checkout * add more tests * test kvcache * kvcache safe * match case * fix segfault due to bad return_softmax * run bench * run seperate for the main functions * just output benchmark * default csv format and time stamp files * non verbsoe bench * Sliding Window Forward (vllm-project#151) * Compress SWA work test case set up debug inputs add fwd ref one mask ref fwd first pass save ref doesnot work for bigger seqlens save new version some causal cases failing found bad cases working new attn new atten works new attn_fwd works reorg n_extra_tokens use seqlen_delta_qk ref fwd works add sliding window to bwd ref test kvcache decode ref work with everything except sliding window add debug code for 12 failing sliding window cases for decode attention_decode_forward_ref_impl mostly works except for alibi fix alibi in attention_decode_forward_ref_impl ref works with normal, varlen & kvcache move stuff around figure out masking old attn inner two inner functions remove load_fn do Lk - Lq like ref unify IS_CAUSAL code in epilogue clean up add args rm inference stuff simplify compute_masking simpler compute mask stub out returning front masking variables remove pointer pass compute ptrs inloop compute block min and max window stub inside inner mask loop trying to use attn_fwd_mask causes issues fix compiler bug when front masking gen specifc types add sliding window and debug statements use identity for v add more taste cases add comments save use k_max_token for clarity disable debug configs basic NON-CAUSAL SLIDING WINDOW non causal sliding window works on the all the shapes non sliding window working in fwd clean up fused bwd seperate old fwd_prefill move configs to utils.py * fix bwd ref bug * skip local cases so that fa output * no sliding window causal green * add backward test skip for sliding window * clean reduce in fwd_kvcache. no is_CASUAL branching * add kvcache masking * kvcache working * fix some bugs in test.py * clean up * Fix Device Segfault (vllm-project#152) * Compress segfault work fix backward segfault rework offset ignore .profile ignore .analysis save * assert the kernel launch device and tensor devices are the same * fix failing asserts * add asserts to fwd * Fix SDMASK bug * Log triton, torch and fa version * Fix fp8 import issues * fix docs (vllm-project#154) * Sliding Window block classification logic (vllm-project#155) * add aiter code * remove aiter stuff * sliding window non causal masking works * causal and sliding window block masking * extract common * clean up typo * helper for swa * ignore .amd * fix last block bug * Enable FA V3 (vllm-project#157) * Compress PA work narrow pa test ref works on most cases inplace ref with new_kv inplace paged attention add pa ref save pa basic paged works save fix swa + causal in pa. Also new_kv only on pa path passing build fa v3 import interface from fa v3 copy fa tests use v3 api clean up rename to match old test support different head sizes remove fp8 basisc passing v3 cases test_flash_attn_varlen_output v3 working isolate bad case for kvcache case passing save use decode is seqused/ cacheseql is given use decode if not varlen basci kvcache v3 working kvcache enable more cases detect kvcache case if seqused_q is non and sequese_k is not None skip failing test find fp8 failing case mha fp8 works fix fp8 MQA/GQA bug clean up more clean up clean up more don't need fp8 dead code remove train code with fp8 stuff fp8 working in kvcache paged + fp8 seems to be working new_kv allowed * clean up * skip hopper race test * clean up more * fix paged + alibi * similar inner paged api * unify _attn_fwd_inner * AITER integration (vllm-project#159) * clean up v2 interface * assert fp8 scale shapes * rotary working * move rotary to impl layers * remove einops * enable rotarry in v3 * create interface * fix descale assert * unify bwd * lint from aiter * clean fp8 api * add api change * assert shapes for v2 * remove ref and bench.py * remove metadata class and clean up * bwd_prefill * one bwd.py * rename * lint * add bwd_change (vllm-project#156) * Tune FP8 Perf (vllm-project#160) * check cu count for gfx942 * create get_cu_count * update repo root * update forward tune * clean up load * use float8_e4m3fnuz * save * show bwd mode * recommend fp8 * use torch.float32 for fp8 kernel * add both best fp16 and fp8 config * tune fp8 backward * descale factors should be b, hk * fp8 bwd working on all primus configs * tune bwd configs * fa v3 tests passing * better warning * clean up bwd launcher * v3 passing * tune more * improve perf * clean up * lint * clean * start tuning gfx950 * tune non causal path * fix bug * save * Skip configs where BLOCK_M2 % BLOCK_N2 != 0 * skip more * stop tuning * fix varlen bug * fix dropout & causal/swa segfault * update the to machine new changes * save * fix more bugs * remove random seed * clean up * update readme * print tensor stats for debug * disable sliding window tests * add rdna configs * fix k partial bug * fix block_size_n bug * fix type check bug --------- Co-authored-by: Aliasger Zaidy <aliasger.zaidy@amd.com> Co-authored-by: Tianxing Wu <tianxing.wu@amd.com>

cleonard530 · 2026-06-24T13:26:57Z

IIUC, are you actually porting our changes to FA3 (in our fork) to the stable API implementation (from upstream code)?

@Harry-Chen Yes, the stable API that was created in upstream got out of sync with the current changes in vLLM's FA3 API, so this PR is to get them back in sync.

Harry-Chen · 2026-06-24T14:02:14Z

IIUC, are you actually porting our changes to FA3 (in our fork) to the stable API implementation (from upstream code)?

@Harry-Chen Yes, the stable API that was created in upstream got out of sync with the current changes in vLLM's FA3 API, so this PR is to get them back in sync.

Great. Could you also create a PR in the vllm repo to change the ref? (I'm not sure whether it could fetch a commit from a PR).

cleonard530 · 2026-06-24T15:23:54Z

IIUC, are you actually porting our changes to FA3 (in our fork) to the stable API implementation (from upstream code)?

@Harry-Chen Yes, the stable API that was created in upstream got out of sync with the current changes in vLLM's FA3 API, so this PR is to get them back in sync.

Great. Could you also create a PR in the vllm repo to change the ref? (I'm not sure whether it could fetch a commit from a PR).

@Harry-Chen Once this commit lands and I update the vllm-project/flash-attention CMakeLists.txt to build with the Stable ABI, the only change needed in vllm-project/vllm should be updating the GIT_TAG in vllm/cmake/external_projects/vllm_flash_attn.cmake

  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
          GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

My original plan was:

PR 1: Sync with the unstable API.
PR 2: Update the build to use the Stable ABI.
PR 3 (in vllm-project/vllm): Update the GIT_TAG to the new flash-attention revision.

I can combine the two PRs in this repository if that would be easier. However, I think it would be preferable for vllm to reference a commit that has already landed on the main branch here, which means this PR would need to merge before the corresponding vllm change.

If we'd like to validate everything in vllm before merging this PR, I could also open a vllm PR that points directly to this branch and use that to run CI. After verification, we could merge this PR, update the vllm PR to reference the merged commit on main, and then merge the vllm change.

What do you think about this?

Hopefully this makes sense, but please let me know if you need me to clear anything up.

Harry-Chen · 2026-06-24T15:42:41Z

If we'd like to validate everything in vllm before merging this PR, I could also open a vllm PR that points directly to this branch and use that to run CI.

Yes this is actually what I meant. I do not have permission in this repo, so I need to find other maintainers and show "hey this is compiling and working as expected" to get it merged, then in the main repo.

I think one PR in this repo is totally fine, since there should not be other users other than vllm itself.

…in static_switch to change TORCH_CHECK to STD_TORCH_CHECK because some files in the binary were using it and it caused one unstable symbol Signed-off-by: Chris Leonard <chleonar@redhat.com>

Harry-Chen

LGTM

Incorporates vllm-project#160 (revert of vllm-project#152, the torch stable-ABI port that broke FLASH_ATTN_MLA_SPARSE). Net change: hopper stable-ABI files only; no cute-DSL / FP8 changes.

cleonard530 · 2026-07-13T12:34:31Z

@@ -2,11 +2,7 @@
 * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.


This commit was created just to show the diff between flash_api.cpp (old) and flash_api_stable.cpp (updates)

cleonard530 force-pushed the update_build_to_torch_stable_abi branch 2 times, most recently from 61d7af3 to c993fa0 Compare June 18, 2026 20:57

cleonard530 changed the title ~~Replace unstable FA3 API with stable ABI port~~ Update stable FA3 API (flash_api_stable.cpp) to be inline with the unstable ABI port (flash_api.cpp) Jun 22, 2026

cleonard530 commented Jun 22, 2026

View reviewed changes

cleonard530 commented Jun 23, 2026

View reviewed changes

cleonard530 added 9 commits June 23, 2026 21:19

made small updates to flash_api_stable.cpp to align better with flash…

18cc90d

…_api.cpp. Returned the CMakeLists.txt and static_switch.h files back to main, will migrate them in a different PR. Signed-off-by: Chris Leonard <chleonar@redhat.com>

copied flash_api_stable.cpp to flash_api.cpp so that the diff is easy…

7d64c73

… to see in github. Will move back in next commit Signed-off-by: Chris Leonard <chleonar@redhat.com>

moving flash_api.cpp back to main

77cd20e

Signed-off-by: Chris Leonard <chleonar@redhat.com>

more commits to make flash_api_stable.cpp align more with flash_abi.cpp

2ece6b3

Signed-off-by: Chris Leonard <chleonar@redhat.com>

updated the stable run_mha_fwd block to align with the unstable api file

ae79a65

Signed-off-by: Chris Leonard <chleonar@redhat.com>

[2nd time] copied flash_api_stable.cpp to flash_api.cpp so that the d…

7da1d8c

…iff is easy to see in github. Will move back in next commit Signed-off-by: Chris Leonard <chleonar@redhat.com>

[2nd time] moving flash_api.cpp back to main

28149df

Signed-off-by: Chris Leonard <chleonar@redhat.com>

removed default arguments to align with legacy stable_abi.cpp

9252523

Signed-off-by: Chris Leonard <chleonar@redhat.com>

cleonard530 force-pushed the update_build_to_torch_stable_abi branch from c827b79 to 9252523 Compare June 23, 2026 21:19

cleonard530 marked this pull request as ready for review June 23, 2026 21:21

update FA3 to build with flash_api_stable. Also, made a small update …

eef3f40

…in static_switch to change TORCH_CHECK to STD_TORCH_CHECK because some files in the binary were using it and it caused one unstable symbol Signed-off-by: Chris Leonard <chleonar@redhat.com>

cleonard530 mentioned this pull request Jun 24, 2026

[Build] Update vllm to point to vllm-project/flash-attention commit that builds FA3 with torch stable API. vllm-project/vllm#46644

Merged

4 tasks

Harry-Chen approved these changes Jun 25, 2026

View reviewed changes

WoosukKwon approved these changes Jun 25, 2026

View reviewed changes

Merge branch 'main' into update_build_to_torch_stable_abi

685018b

WoosukKwon merged commit b3964b1 into vllm-project:main Jun 25, 2026
1 check passed

MatthewBonanni mentioned this pull request Jul 7, 2026

Sync upstream #149

Merged

This was referenced Jul 10, 2026

Revert "Update stable FA3 API (flash_api_stable.cpp) to be inline with the unstable ABI port (flash_api.cpp)" #160

Merged

[Revert] [Build] Update vllm ...builds FA3 with torch stable API vllm-project/vllm#48269

Merged

elvircrn mentioned this pull request Jul 11, 2026

fix: align get_scheduler_metadata with mha_fwd in flash_api.cpp #162

Draft

cleonard530 commented Jul 13, 2026

View reviewed changes

cleonard530 mentioned this pull request Jul 16, 2026

Re-migrate FA3 hopper API to the PyTorch stable ABI #165

Merged

		@@ -2,11 +2,7 @@
		* Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.

Uh oh!

Conversation

cleonard530 commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

vLLM Test Plan/Result

Test Result

Uh oh!

cleonard530 Jun 22, 2026

Choose a reason for hiding this comment

Uh oh!

cleonard530 Jun 22, 2026

Choose a reason for hiding this comment

Uh oh!

cleonard530 Jun 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cleonard530 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

cleonard530 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

cleonard530 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

cleonard530 Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

Harry-Chen commented Jun 24, 2026

Uh oh!

cleonard530 commented Jun 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Harry-Chen commented Jun 24, 2026

Uh oh!

cleonard530 commented Jun 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Harry-Chen commented Jun 24, 2026

Uh oh!

Harry-Chen left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cleonard530 Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cleonard530 commented Jun 18, 2026 •

edited

Loading

cleonard530 Jun 22, 2026 •

edited

Loading

cleonard530 commented Jun 24, 2026 •

edited

Loading

cleonard530 commented Jun 24, 2026 •

edited

Loading