Re-organize SLL ops, pt 5

Summary: - Re-organize `jagged_jagged_bmm` and `jagged_dense_bmm` Differential Revision: D68935929
q10 · Feb 1, 2025 · 7d0c24d · 7d0c24d
1 parent b9ef10e
commit 7d0c24d
Show file tree

Hide file tree

Showing 4 changed files with 434 additions and 417 deletions.
diff --git a/fbgemm_gpu/fbgemm_gpu/sll/__init__.py b/fbgemm_gpu/fbgemm_gpu/sll/__init__.py
@@ -37,9 +37,8 @@
     array_jagged_bmm_jagged_out,
     dense_jagged_cat_jagged_out,
     jagged2_to_padded_dense,
-    jagged_dense_bmm,
+    # jagged_dense_bmm,
     jagged_dense_elementwise_mul_jagged_out,
-    jagged_jagged_bmm,
     jagged_jagged_bmm_jagged_out,
     triton_jagged_self_substraction_jagged_out,
 )
@@ -270,14 +269,6 @@
 
 # pyre-ignore[5]
 sll_gpu_registrations = {
-    "sll_jagged_dense_bmm": {
-        "CUDA": jagged_dense_bmm,
-        "AutogradCUDA": jagged_dense_bmm,
-    },
-    "sll_jagged_jagged_bmm": {
-        "CUDA": jagged_jagged_bmm,
-        "AutogradCUDA": jagged_jagged_bmm,
-    },
     "sll_dense_jagged_cat_jagged_out": {
         "CUDA": dense_jagged_cat_jagged_out,
     },

diff --git a/fbgemm_gpu/fbgemm_gpu/sll/triton/__init__.py b/fbgemm_gpu/fbgemm_gpu/sll/triton/__init__.py
@@ -7,6 +7,13 @@
 
 # pyre-strict
 
+from fbgemm_gpu.sll.triton.triton_jagged_bmm import (  # noqa F401
+    jagged_dense_bmm,
+    jagged_jagged_bmm,
+    JaggedDenseBmm,  # noqa F401
+    JaggedJaggedBmm,  # noqa F401
+)
+
 from fbgemm_gpu.sll.triton.triton_jagged_dense_elementwise_add import (  # noqa F401
     jagged_dense_elementwise_add,
     JaggedDenseAdd,  # noqa F401
@@ -36,6 +43,14 @@
 
 # pyre-ignore[5]
 op_registrations = {
+    "sll_jagged_dense_bmm": {
+        "CUDA": jagged_dense_bmm,
+        "AutogradCUDA": jagged_dense_bmm,
+    },
+    "sll_jagged_jagged_bmm": {
+        "CUDA": jagged_jagged_bmm,
+        "AutogradCUDA": jagged_jagged_bmm,
+    },
     "sll_jagged_softmax": {
         "CUDA": jagged_softmax,
         "AutogradCUDA": jagged_softmax,