Re-organize SLL ops, pt 4 (pytorch#3644)

Summary: X-link: facebookresearch/FBGEMM#720 Pull Request resolved: pytorch#3644 - Re-organize `jagged_flash_attention_basic`, `jagged_softmax`, and `jagged2_softmax` Differential Revision: D68924000
q10 · Jan 31, 2025 · f601628 · f601628
1 parent e7312a3
commit f601628
Show file tree

Hide file tree

Showing 5 changed files with 1,156 additions and 1,121 deletions.
diff --git a/fbgemm_gpu/fbgemm_gpu/sll/__init__.py b/fbgemm_gpu/fbgemm_gpu/sll/__init__.py
@@ -36,14 +36,11 @@
 from fbgemm_gpu.sll.triton_sll import (  # noqa F401
     array_jagged_bmm_jagged_out,
     dense_jagged_cat_jagged_out,
-    jagged2_softmax,
     jagged2_to_padded_dense,
     jagged_dense_bmm,
     jagged_dense_elementwise_mul_jagged_out,
-    jagged_flash_attention_basic,
     jagged_jagged_bmm,
     jagged_jagged_bmm_jagged_out,
-    jagged_softmax,
     triton_jagged_self_substraction_jagged_out,
 )
 
@@ -295,14 +292,6 @@
         "CUDA": jagged_dense_elementwise_mul_jagged_out,
         "AutogradCUDA": jagged_dense_elementwise_mul_jagged_out,
     },
-    "sll_jagged_softmax": {
-        "CUDA": jagged_softmax,
-        "AutogradCUDA": jagged_softmax,
-    },
-    "sll_jagged2_softmax": {
-        "CUDA": jagged2_softmax,
-        "AutogradCUDA": jagged2_softmax,
-    },
     "sll_array_jagged_bmm_jagged_out": {
         "CUDA": array_jagged_bmm_jagged_out,
         "AutogradCUDA": array_jagged_bmm_jagged_out,
@@ -311,10 +300,6 @@
         "CUDA": jagged_jagged_bmm_jagged_out,
         "AutogradCUDA": jagged_jagged_bmm_jagged_out,
     },
-    "sll_jagged_flash_attention_basic": {
-        "CUDA": jagged_flash_attention_basic,
-        "AutogradCUDA": jagged_flash_attention_basic,
-    },
 }
 
 for op_name, dispatches in sll_cpu_registrations.items():

diff --git a/fbgemm_gpu/fbgemm_gpu/sll/triton/__init__.py b/fbgemm_gpu/fbgemm_gpu/sll/triton/__init__.py
@@ -7,23 +7,43 @@
 
 # pyre-strict
 
-
 from fbgemm_gpu.sll.triton.triton_jagged_dense_elementwise_add import (  # noqa F401
     jagged_dense_elementwise_add,
     JaggedDenseAdd,  # noqa F401
 )
+
 from fbgemm_gpu.sll.triton.triton_jagged_dense_flash_attention import (  # noqa F401
     jagged_dense_flash_attention,
     JaggedDenseFlashAttention,  # noqa F401
 )
 
+from fbgemm_gpu.sll.triton.triton_jagged_flash_attention_basic import (  # noqa F401
+    jagged_flash_attention_basic,
+    JaggedFlashAttentionBasic,  # noqa F401
+)
+
+from fbgemm_gpu.sll.triton.triton_jagged_softmax import (  # noqa F401
+    jagged2_softmax,
+    Jagged2Softmax,  # noqa F401
+    jagged_softmax,
+    JaggedSoftmax,  # noqa F401
+)
+
 from fbgemm_gpu.sll.triton.triton_multi_head_jagged_flash_attention import (  # noqa F401
     multi_head_jagged_flash_attention,
     MultiHeadJaggedFlashAttention,  # noqa F401
 )
 
 # pyre-ignore[5]
 op_registrations = {
+    "sll_jagged_softmax": {
+        "CUDA": jagged_softmax,
+        "AutogradCUDA": jagged_softmax,
+    },
+    "sll_jagged2_softmax": {
+        "CUDA": jagged2_softmax,
+        "AutogradCUDA": jagged2_softmax,
+    },
     "sll_jagged_dense_elementwise_add": {
         "CUDA": jagged_dense_elementwise_add,
         "AutogradCUDA": jagged_dense_elementwise_add,
@@ -32,6 +52,10 @@
         "CUDA": jagged_dense_flash_attention,
         "AutogradCUDA": jagged_dense_flash_attention,
     },
+    "sll_jagged_flash_attention_basic": {
+        "CUDA": jagged_flash_attention_basic,
+        "AutogradCUDA": jagged_flash_attention_basic,
+    },
     "sll_multi_head_jagged_flash_attention": {
         "CUDA": multi_head_jagged_flash_attention,
         "AutogradCUDA": multi_head_jagged_flash_attention,