65 Commits

Author SHA1 Message Date
skarupke bc0e4ac014 Fix use-after-free in FA3 deterministic mode. The pytorch caching allocator actually saves us here, but if you turn it off, then compute-sanitizer will detect this. (#2063) 2025-12-12 06:47:31 -08:00
Guilherme Leobas e1944ba9cb Don't return mutated variables in mha_bwd 2025-11-17 19:17:52 +00:00
Guilherme Leobas c7697bbf3e Add torch.compile support to flash attention 3 2025-11-17 19:17:52 +00:00
Johnny add175637c C++11 fix warnings (#1904)
* errors are with C++11 narrowing warnings (treated as errors in strict builds) when initializing at::cuda::CUDAGuard with a non-constant char cast to c10::DeviceIndex (signed char).

* errors are with C++11 narrowing warnings (treated as errors in strict builds) when initializing at::cuda::CUDAGuard with a non-constant char cast to c10::DeviceIndex (signed char).

* errors are with C++11 narrowing warnings (treated as errors in strict builds) when initializing at::cuda::CUDAGuard with a non-constant char cast to c10::DeviceIndex (signed char).

* errors are with C++11 narrowing warnings (treated as errors in strict builds) when initializing at::cuda::CUDAGuard with a non-constant char cast to c10::DeviceIndex (signed char).

* Update flash_api_stable.cpp

* upstream cutlass v4.2.1 csrc
2025-09-25 01:22:47 -07:00
jayhshah 5c1627a7a1 Improve causal backward determinism perf with SPT schedule (#1893)
* add spt scheduler for causal bwd determinism

* add new torch check for det hdim 256 to stable api
2025-09-17 14:58:45 -07:00
y-sq 832d5448ce Update the initialization of dk/dv_semaphore (#1839)
When testing the deterministic option for the GQA case, we found it fell into deadlock issues. Initialization dk and dv_semaphore to zeros to fix this issue.
2025-08-25 07:44:22 -04:00
jayhshah 199401d31f Add sorting and head swizzle to varlen scheduler (#1823)
* use LPT order in varlen kernel

* add prefill decode benchmark script

* add sort in prepare

* add full implementation:

* add varlen kvhead swizzle

* add settings for swizzle ablation

* add correction term for sort when causal

* remove ablation options from frontend and clean up comments

* add comments in prepare kernel

* remove debug code and scripts

* put back defaults in tests

* remove excess Nones returned in python interface for varlen

* revert opinionated change to setup.py on cuda version 12.9

* force inline sort op and make east const

* more templating in varlen scheduler to cure some register spilling

* fix exploding build by splitting compilation and add qol macros for hdimdiff

* fix metadata mismatch with seqlenk in test script

* extend prepare kernel to >992 batches and always call it for varlen

* do inter-batch sort per every 992 batches

* better names in combine and fix prepare condition in api
2025-08-21 19:44:03 -07:00
dan_the_3rd 847025a002 [fa3] API default values + backward compatibility (#1700)
* Make Flash3 backward compatible

* Finish tests
2025-06-07 21:59:22 -07:00
dan_the_3rd fc27c4f6aa [fa3] Some fixes for windows build (#1698)
* switch fw

* fix

* f

* f

* f

* constexpr conds on arch

* format

* format again
2025-06-07 21:58:49 -07:00
Jane (Yuan) Xu 0e79d71175 [BE] use more minimal torch headers for hopper/flash_api.cpp (#1674)
Co-authored-by: Tri Dao <tridao@users.noreply.github.com>
2025-05-21 21:44:03 -07:00
dan_the_3rd db4baba2ca [fa3] Use Python stable ABI (#1662)
* Use Python stable ABI

* Remove useless macro

* Add 'py_limited_api=True'

* Default value for 'num_splits'

* Update hopper/flash_api.cpp

Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com>

---------

Co-authored-by: Tri Dao <tridao@users.noreply.github.com>
Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com>
2025-05-21 21:42:37 -07:00
Sanghun Cho 37c816ab0d Support hdimQK != hdimV backward (#1604)
* separate d & dv (interface)

* separate d & dv (api)

* separate d & dv (template)

* separate d & dv (mainloop)

* separate d & dv (epilogue)

* update test

* disable backward test when attention_chunk != 0

* extend backward d > dv to d != dv

---------

Co-authored-by: monk.ey <monk.ey@kakaocorp.com>
2025-04-23 20:17:27 -07:00
Tri Dao 757c5ad577 Fix kvcache tests for attention_chunk when precomputing metadata 2025-04-12 01:32:51 -04:00
Tri Dao 7bb8e8249d Make attention_chunk work for non-causal cases 2025-04-11 23:56:38 -04:00
Tri Dao 909eb7ce7c Implement attention_chunk 2025-04-10 04:51:52 -04:00
Tri Dao 1a58058a6d Fix zeroing out the scheduler semaphore when reusing metadata 2025-03-29 01:29:39 -04:00
Tri Dao 29ef580560 Adjust warp scheduler sync for HasQv case 2025-03-22 17:46:12 -04:00
Tri Dao f6a294a244 Update benchmark with GLA 2025-03-22 17:15:29 -04:00
Tri Dao df11fcae26 Support hdim 64,256 2025-03-22 16:22:07 -04:00
Tri Dao 7ae5f8c8fe Add option for rotary_seqlens 2025-03-21 02:24:16 -07:00
Tri Dao fa60e7cc97 Add option to precompute scheduler metadata 2025-03-15 16:41:29 -04:00
Tri Dao 897c84539a Fix: num_splits_dynamic_ptr needs to be set before get_num_splits 2025-03-14 00:38:03 -04:00
Tri Dao 46e1d4a1c7 Simplify prepare_varlen_num_blocks_kernel, restrict to batch <= 992 2025-03-13 01:45:05 -04:00
Tri Dao 000090d02f Enable PDL 2025-03-12 20:32:26 -04:00
Tri Dao 430954a8a1 Adjust heuristic for get_pagedkv_tma 2025-03-12 18:40:36 -04:00
Tri Dao 1d30bb4cd3 Enable TMA when page size is a multiple of kBlockN 2025-03-08 22:19:39 -05:00
Tri Dao 81643fa0ea For each batch, if num_splits=1, write to O instead of O_partial 2025-03-08 14:50:54 -05:00
Tri Dao d82bbf2692 Fix causal scheduler not considering hdim_v != hdim 2025-03-05 16:51:48 -05:00
Tri Dao 45c48afb2b Add option for WG1 to use RS MMA but WG2 using SS MMA 2025-03-03 23:54:24 -05:00
Tri Dao 5458c78e6d Remove sink token
It wasn't working anyway
2025-03-01 21:03:47 -05:00
Tri Dao 6752d62aa4 Add dynamic splits 2025-02-24 03:37:05 -05:00
Ted Zadouri 06e34f62d1 Enable MLA flag in FA3 (rope=64, latent=512) (#1504)
* Enable MLA flag in FA3 (rope=64, latent=512)

* updated HasQv in flash_fwd_launch_template.h
2025-02-22 18:24:44 -08:00
Tri Dao b36ad4ef76 Use split for super long sequences that don't fit into L2 2025-02-19 02:09:07 -05:00
Tri Dao 5378bc3204 Tile fwd_combine kernel along headdim, don't need kBlockM > 128 2025-02-08 16:30:51 -05:00
Tri Dao 5fc5ebf82b Fix num_splits heuristic again when PackGQA 2025-02-08 16:21:44 -05:00
Tri Dao 5fab938555 Fix num_splits heuristic being called before get_pack_gqa 2025-02-08 01:25:09 -05:00
Tri Dao 4e8496a781 Support Qv 2025-02-07 22:55:07 -05:00
Tri Dao ed53b5fc4c Fix compilation for FP8 when kHeadDimV != kHeadDim 2025-02-07 22:25:17 -05:00
Tri Dao 2a204125ae Implement HeadDim_V != HeadDim_QK, support hdimQK=192, hdimV=128 2025-02-03 00:19:25 -05:00
Lu Fang 74aed78373 Replace c10::optional with std::optional in flash_attn
Signed-off-by: Lu Fang <lufang@fb.com>
2025-01-13 14:37:27 -08:00
Tri Dao 2ac6c986be Fix Sm80 tile_count_semaphore, adjust test tolerance 2025-01-10 23:44:51 +07:00
Tri Dao e94f7e89dc Always enable PackGQA is Split to reduce compilation and binary size 2025-01-10 19:55:31 +07:00
Tri Dao 40fa35acd8 Always enable PackGQA if PagedKV to reduce compilation and bin size 2025-01-10 17:06:12 +07:00
Tri Dao a84a237d2a Split bwd softcap compilation units for Sm80 2025-01-09 23:57:05 +07:00
Tri Dao 1e3208566a Tune tile sizes for compilation 2025-01-09 17:59:57 +07:00
Tri Dao 84f1287e42 Rename bool_constant<true> to true_type, same w bool_constant<false> 2025-01-09 15:51:33 +07:00
Tri Dao df5fe55264 Change tile sizes for Sm8x to reduce stack frame 2025-01-09 15:10:15 +07:00
Tri Dao 8dd0b479d5 Always enable PackGQA for Sm8x to reduce compilation and binary size 2025-01-09 11:51:21 +07:00
Tri Dao 7bc3f031a4 Compile for both Sm80 and Sm90 2025-01-08 18:16:58 +07:00
Tri Dao 7a802796e1 Big refactor and update 2025-01-07 23:34:37 +07:00