Re-organize SLL ops, pt 3 (pytorch#728)

Summary:
Pull Request resolved: facebookresearch/FBGEMM#728

X-link: pytorch#3652

- Re-organize `jagged_dense_elementwise_add`

Reviewed By: sryap

Differential Revision: D68923208

fbshipit-source-id: 1c3f22d9588f11a664fb02843b457ea33bb40f9c

Author: q10

fbgemm_gpu/fbgemm_gpu/sll/__init__.py

@@ -39,7 +39,6 @@
     jagged2_softmax,
     jagged2_to_padded_dense,
     jagged_dense_bmm,
-    jagged_dense_elementwise_add,
     jagged_dense_elementwise_mul_jagged_out,
     jagged_flash_attention_basic,
     jagged_jagged_bmm,
@@ -316,10 +315,6 @@
         "CUDA": jagged_flash_attention_basic,
         "AutogradCUDA": jagged_flash_attention_basic,
     },
-    "sll_jagged_dense_elementwise_add": {
-        "CUDA": jagged_dense_elementwise_add,
-        "AutogradCUDA": jagged_dense_elementwise_add,
-    },
 }
 
 for op_name, dispatches in sll_cpu_registrations.items():

fbgemm_gpu/fbgemm_gpu/sll/triton/__init__.py

@@ -8,18 +8,26 @@
 # pyre-strict
 
 
-from fbgemm_gpu.sll.triton.jagged_dense_flash_attention import (  # noqa F401
+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.multi_head_jagged_flash_attention import (  # 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_dense_elementwise_add": {
+        "CUDA": jagged_dense_elementwise_add,
+        "AutogradCUDA": jagged_dense_elementwise_add,
+    },
     "sll_jagged_dense_flash_attention": {
         "CUDA": jagged_dense_flash_attention,
         "AutogradCUDA": jagged_dense_flash_attention,

fbgemm_gpu/fbgemm_gpu/sll/triton/triton_jagged_dense_elementwise_add.py

@@ -0,0 +1,52 @@
+# 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
+
+from fbgemm_gpu.triton.jagged.triton_jagged_tensor_ops import (
+    dense_to_jagged,
+    jagged_to_dense,
+)
+
+
+class JaggedDenseAdd(torch.autograd.Function):
+    @staticmethod
+    # pyre-fixme
+    def forward(
+        ctx, x: torch.Tensor, x_offsets: torch.Tensor, y: torch.Tensor, max_seq_len: int
+    ):
+        ctx.save_for_backward(x_offsets)
+        ctx.max_seq_len = max_seq_len
+        # TODO: what should be the correct behavior when jagged values has length > max seq len?
+        # current behavior is to not truncate jagged values
+        # similar for backward grad_output
+        return dense_to_jagged(
+            y, [x_offsets], operation_function="add", operation_jagged_values=x
+        )[0]
+
+    @staticmethod
+    # pyre-fixme
+    def backward(ctx, grad_output: torch.Tensor):
+        (offsets,) = ctx.saved_tensors
+        grad_dense = jagged_to_dense(grad_output, [offsets], [ctx.max_seq_len])
+        return grad_output, None, grad_dense, None
+
+
+def jagged_dense_elementwise_add(
+    x: torch.Tensor,
+    x_offsets: torch.Tensor,
+    y: torch.Tensor,
+    max_seq_len: int,
+    use_fbgemm_kernel: bool = True,
+):
+    if use_fbgemm_kernel:
+        return torch.ops.fbgemm.jagged_dense_elementwise_add_jagged_output(
+            x, [x_offsets], y
+        )[0]
+    else:
+        return JaggedDenseAdd.apply(x, x_offsets, y, max_seq_len)

fbgemm_gpu/fbgemm_gpu/sll/triton/jagged_dense_flash_attention.py renamed to fbgemm_gpu/fbgemm_gpu/sll/triton/triton_jagged_dense_flash_attention.py

@@ -12,11 +12,6 @@
 import triton
 import triton.language as tl
 
-from fbgemm_gpu.triton.jagged.triton_jagged_tensor_ops import (
-    dense_to_jagged,
-    jagged_to_dense,
-)
-
 
 def set_block_size(N: int) -> int:
     if N > 64:
@@ -2591,41 +2586,3 @@ def jagged_flash_attention_basic(
     )
 
     return jagged_O
-
-
-class JaggedDenseAdd(torch.autograd.Function):
-    @staticmethod
-    # pyre-fixme
-    def forward(
-        ctx, x: torch.Tensor, x_offsets: torch.Tensor, y: torch.Tensor, max_seq_len: int
-    ):
-        ctx.save_for_backward(x_offsets)
-        ctx.max_seq_len = max_seq_len
-        # TODO: what should be the correct behavior when jagged values has length > max seq len?
-        # current behavior is to not truncate jagged values
-        # similar for backward grad_output
-        return dense_to_jagged(
-            y, [x_offsets], operation_function="add", operation_jagged_values=x
-        )[0]
-
-    @staticmethod
-    # pyre-fixme
-    def backward(ctx, grad_output: torch.Tensor):
-        (offsets,) = ctx.saved_tensors
-        grad_dense = jagged_to_dense(grad_output, [offsets], [ctx.max_seq_len])
-        return grad_output, None, grad_dense, None
-
-
-def jagged_dense_elementwise_add(
-    x: torch.Tensor,
-    x_offsets: torch.Tensor,
-    y: torch.Tensor,
-    max_seq_len: int,
-    use_fbgemm_kernel: bool = True,
-):
-    if use_fbgemm_kernel:
-        return torch.ops.fbgemm.jagged_dense_elementwise_add_jagged_output(
-            x, [x_offsets], y
-        )[0]
-    else:
-        return JaggedDenseAdd.apply(x, x_offsets, y, max_seq_len)

fbgemm_gpu/fbgemm_gpu/sll/triton/multi_head_jagged_flash_attention.py renamed to fbgemm_gpu/fbgemm_gpu/sll/triton/triton_multi_head_jagged_flash_attention.py

@@ -9,9 +9,10 @@
 
 import unittest
 
+import fbgemm_gpu.sll  # noqa F401
 import hypothesis.strategies as st
 import torch
-from fbgemm_gpu.sll.triton_sll import jagged_dense_elementwise_add  # noqa
+
 from hypothesis import given, settings
 
 from .common import open_source