add persistent variant (#77)

Summary: Hongtao identified the performance issue with the initial implementation and updated the assignments of tiles to each SM. Performance with warp specialization (Batch, Heads, SeqLen, Dhead) triton_tutorial_flash_v2_tma_ws_persistent-tflops triton_tutorial_flash_v2_tma_ws-tflops triton_tutorial_flash_v2-tflops ------------------------------- --------------------------------------------------- ---------------------------------------- --------------------------------- (8, 16, 8192, 128) 516.164 490.451 423.905 Pull Request resolved: #77 Reviewed By: xuzhao9, htyu Differential Revision: D66463179 Pulled By: manman-ren fbshipit-source-id: 14fecc1a1449828bfd82600bd161596349da3084
pytorch-labs · Dec 2, 2024 · 0a82d3d · 0a82d3d
1 parent 7e1f269
commit 0a82d3d
Show file tree

Hide file tree

Showing 4 changed files with 306 additions and 10 deletions.
diff --git a/test/test_gpu/skip_tests_h100_pytorch.yaml b/test/test_gpu/skip_tests_h100_pytorch.yaml
@@ -16,6 +16,7 @@ flash_attention:
   - triton_tutorial_flash_v2_tma
   - triton_tutorial_flash_v2_ws
   - triton_tutorial_flash_v2_tma_ws
+  - triton_tutorial_flash_v2_tma_ws_persistent
 fp8_attention:
   - colfax_fmha
   # triton_flash_v2 requires triton-main

diff --git a/test/test_gpu/skip_tests_h100_triton_main.yaml b/test/test_gpu/skip_tests_h100_triton_main.yaml
@@ -10,6 +10,7 @@ flash_attention:
   # _ws kernels require Triton with warp specialization
   - triton_tutorial_flash_v2_ws
   - triton_tutorial_flash_v2_tma_ws
+  - triton_tutorial_flash_v2_tma_ws_persistent
 fp8_attention:
   # fb-only kernel
   - colfax_fmha