Re-organize SLL ops, pt 1 (pytorch#3642)

Summary:
Pull Request resolved: pytorch#3642

X-link: https://github.com/facebookresearch/FBGEMM/pull/718

- Re-organize SLL ops, pt 1

Reviewed By: sryap

Differential Revision: D68915217

fbshipit-source-id: 6208ef53c1740c5dbf89534cc2301c10242b82ea

Author: q10

fbgemm_gpu/fbgemm_gpu/sll/__init__.py

@@ -46,7 +46,6 @@
     jagged_jagged_bmm,
     jagged_jagged_bmm_jagged_out,
     jagged_softmax,
-    multi_head_jagged_flash_attention,
     triton_jagged_self_substraction_jagged_out,
 )
 
@@ -326,15 +325,16 @@
         "CUDA": jagged_dense_flash_attention,
         "AutogradCUDA": jagged_dense_flash_attention,
     },
-    "sll_multi_head_jagged_flash_attention": {
-        "CUDA": multi_head_jagged_flash_attention,
-        "AutogradCUDA": multi_head_jagged_flash_attention,
-    },
 }
 
 for op_name, dispatches in sll_cpu_registrations.items():
     lib.register(op_name, dispatches)
 
 if torch.cuda.is_available():
+    from fbgemm_gpu.sll.triton import op_registrations
+
+    for op_name, dispatches in op_registrations.items():
+        lib.register(op_name, dispatches)
+
     for op_name, dispatches in sll_gpu_registrations.items():
         lib.register(op_name, dispatches)

fbgemm_gpu/fbgemm_gpu/sll/triton/__init__.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+
+from fbgemm_gpu.sll.triton.multi_head_jagged_flash_attention import (  # noqa F401
+    multi_head_jagged_flash_attention,
+    MultiHeadJaggedFlashAttention,
+)
+
+op_registrations = {
+    "sll_multi_head_jagged_flash_attention": {
+        "CUDA": multi_head_jagged_flash_attention,
+        "AutogradCUDA": multi_head_jagged_flash_attention,
+    },
+}

fbgemm_gpu/fbgemm_gpu/sll/triton/common.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+import torch
+
+
+def expect_contiguous(x: torch.Tensor) -> torch.Tensor:
+    if not x.is_contiguous():
+        return x.contiguous()
+    else:
+        return x