fbgemm-gpu-genai-nightly 2025.12.19__cp310-cp310-manylinux_2_28_x86_64.whl

fbgemm_gpu/sll/meta/__init__.py

@@ -0,0 +1,35 @@
+#!/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.meta.meta_sll import (  # noqa F401
+    meta_array_jagged_bmm_jagged_out,
+    meta_jagged2_softmax,
+    meta_jagged_dense_elementwise_mul_jagged_out,
+    meta_jagged_jagged_bmm_jagged_out,
+    meta_jagged_self_substraction_jagged_out,
+)
+
+# pyre-ignore[5]
+op_registrations = {
+    "sll_jagged_self_substraction_jagged_out": {
+        "Meta": meta_jagged_self_substraction_jagged_out,
+    },
+    "sll_jagged_dense_elementwise_mul_jagged_out": {
+        "Meta": meta_jagged_dense_elementwise_mul_jagged_out,
+    },
+    "sll_jagged2_softmax": {
+        "AutogradMeta": meta_jagged2_softmax,
+    },
+    "sll_array_jagged_bmm_jagged_out": {
+        "AutogradMeta": meta_array_jagged_bmm_jagged_out,
+    },
+    "sll_jagged_jagged_bmm_jagged_out": {
+        "AutogradMeta": meta_jagged_jagged_bmm_jagged_out,
+    },
+}

fbgemm_gpu/sll/meta/meta_sll.py

@@ -0,0 +1,337 @@
+# 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
+import torch
+
+
+def meta_jagged_self_substraction_jagged_out(
+    jagged_A: torch.Tensor,
+    offsets_a: torch.Tensor,
+    offsets_b: torch.Tensor,
+    max_seq_len: int,
+) -> torch.Tensor:
+    return torch.empty(
+        [torch.library.get_ctx().new_dynamic_size()],
+        dtype=jagged_A.dtype,
+        device=jagged_A.device,
+    )
+
+
+class MetaJaggedDenseElementwiseMul(torch.autograd.Function):
+    @staticmethod
+    # pyre-fixme
+    def forward(
+        ctx,  # pyre-ignore [2]
+        x: torch.Tensor,
+        y: torch.Tensor,
+        x_seq_lengths: torch.Tensor,
+        x_offsets: torch.Tensor,
+        max_seq_len: int,
+    ) -> torch.Tensor:
+        ctx.max_seq_len = max_seq_len
+
+        ctx.save_for_backward(
+            x,
+            y,
+            x_seq_lengths,
+            x_offsets,
+        )
+
+        total_L = x.size(0)
+        jagged_C = torch.zeros((total_L), device=x.device, dtype=x.dtype)
+
+        return jagged_C
+
+    @staticmethod
+    # pyre-fixme
+    def backward(ctx, grad_output: torch.Tensor):
+        (
+            x,
+            y,
+            x_seq_lengths,
+            x_offsets,
+        ) = ctx.saved_tensors
+
+        total_L = grad_output.size(0)
+        jagged_C = torch.zeros(
+            (total_L), device=grad_output.device, dtype=grad_output.dtype
+        )
+
+        return jagged_C, None, None, None, None
+
+
+def meta_jagged_dense_elementwise_mul_jagged_out(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    x_seq_lengths: torch.Tensor,
+    x_offsets: torch.Tensor,
+    max_seq_len: int,
+) -> torch.Tensor:
+    return MetaJaggedDenseElementwiseMul.apply(
+        x,
+        y,
+        x_seq_lengths,
+        x_offsets,
+        max_seq_len,
+    )
+
+
+class Jagged2SoftmaxMeta(torch.autograd.Function):
+    @staticmethod
+    # pyre-fixme
+    def forward(
+        # pyre-fixme[2]: Parameter must be annotated.
+        ctx,
+        x: torch.Tensor,
+        x_offsets: torch.Tensor,
+        row_offsets: torch.Tensor,
+        head_offsets: torch.Tensor,
+        max_seq_len_row: int,
+        max_seq_len_head: int,
+        transpose: bool = True,
+    ) -> torch.Tensor:
+        y = torch.rand(x.size(0), device=x.device, dtype=x.dtype)
+
+        ctx.save_for_backward(y, x_offsets, row_offsets, head_offsets)
+        ctx.max_seq_len_row = max_seq_len_row
+        ctx.max_seq_len_head = max_seq_len_head
+
+        return y
+
+    @staticmethod
+    # pyre-fixme
+    def backward(ctx, grad_output: torch.Tensor):
+        y, x_offsets, row_offsets, head_offsets = ctx.saved_tensors
+        grad = torch.rand(y.size(), device=y.device, dtype=y.dtype)
+
+        return grad, None, None, None, None, None, None
+
+
+def meta_jagged2_softmax(
+    x: torch.Tensor,
+    offsets: torch.Tensor,
+    offsets_total: torch.Tensor,
+    max_seq_len: int,
+    transpose: bool,
+) -> torch.Tensor:
+    """
+    Meta version of jagged2 softmax: [sum(softmax([B_i, B_i]))]
+    """
+    return Jagged2SoftmaxMeta.apply(
+        x,
+        offsets_total,
+        offsets,
+        offsets,
+        max_seq_len,
+        max_seq_len,
+        transpose,
+    )
+
+
+class ArrayJaggedBmmNopadding(torch.autograd.Function):
+    """
+    Compute batch matrix multiplication between JaggedTensor and JaggedTensor without padding.
+    z = X * Y
+    x: [Sum_B(N_i, N_i)]
+    y: [sum_B(N_i), D]
+    z: [sum_B(N_i), D]
+    """
+
+    @staticmethod
+    # pyre-fixme
+    def forward(
+        # pyre-fixme[2]: Parameter must be annotated.
+        ctx,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        x_lengths: torch.Tensor,
+        x_offsets: torch.Tensor,
+        y_lengths: torch.Tensor,
+        y_offsets: torch.Tensor,
+        z_lengths: torch.Tensor,
+        z_offsets: torch.Tensor,
+        max_seq_len: int,
+        # pyre-fixme[2]: Parameter must be annotated.
+        allow_tf32,
+    ) -> torch.Tensor:
+        ctx.allow_tf32 = allow_tf32
+        ctx.max_seq_len = max_seq_len
+
+        ctx.save_for_backward(
+            x,
+            y,
+            x_lengths,
+            y_lengths,
+            z_lengths,
+            x_offsets,
+            y_offsets,
+            z_offsets,
+        )
+
+        D = y.size(1)
+        L = y.size(0)
+        # gradients of the emb vectors beyond max_seq_len is set to zeros
+        jagged_C = torch.zeros((L, D), device=y.device, dtype=y.dtype)
+        return jagged_C
+
+    @staticmethod
+    # pyre-fixme
+    def backward(ctx, grad_output: torch.Tensor):
+        """
+        z = X * Y
+        dX = dZ * YT
+        dY = XT * dZ
+
+        dZ: [sum_B(N_i), D]
+        YT: [D, sum_B(N_i)] call Y.T
+        XT: transposed
+        Z: [sum_B(N_i), D]
+        """
+
+        (
+            x,
+            y,
+            x_lengths,
+            y_lengths,
+            z_lengths,
+            x_offsets,
+            y_offsets,
+            z_offsets,
+        ) = ctx.saved_tensors
+
+        grad_x = torch.zeros(
+            (x.size()), device=grad_output.device, dtype=grad_output.dtype
+        )
+
+        # gradients of the emb vectors beyond max_seq_len is set to zeros
+        grad_y = torch.zeros(
+            y.size(), device=grad_output.device, dtype=grad_output.dtype
+        )
+        return (
+            grad_x,
+            grad_y,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+# pyre-fixme[3]: Return type must be annotated.
+def meta_array_jagged_bmm_jagged_out(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    x_lengths: torch.Tensor,
+    x_offsets: torch.Tensor,
+    y_lengths: torch.Tensor,
+    y_offsets: torch.Tensor,
+    z_lengths: torch.Tensor,
+    z_offsets: torch.Tensor,
+    max_seq_len: int,
+    allow_tf32: bool = True,
+):
+    return ArrayJaggedBmmNopadding.apply(
+        x,
+        y,
+        x_lengths,
+        x_offsets,
+        y_lengths,
+        y_offsets,
+        z_lengths,
+        z_offsets,
+        max_seq_len,
+        allow_tf32,
+    )
+
+
+class JaggedJaggedBmmNoPaddingMeta(torch.autograd.Function):
+    @staticmethod
+    # pyre-fixme
+    def forward(
+        # pyre-fixme[2]: Parameter must be annotated.
+        ctx,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        x_lengths: torch.Tensor,
+        x_offsets: torch.Tensor,
+        y_lengths: torch.Tensor,
+        y_offsets: torch.Tensor,
+        z_lengths: torch.Tensor,
+        z_offsets: torch.Tensor,
+        max_seq_len: int,
+        # pyre-fixme[2]: Parameter must be annotated.
+        allow_tf32,
+    ):
+        assert x.size(1) == y.size(0), "incompatible dimensions"
+
+        ctx.allow_tf32 = allow_tf32
+        ctx.max_seq_len = max_seq_len
+
+        ctx.save_for_backward(
+            x,
+            y,
+            x_lengths,
+            y_lengths,
+            z_lengths,
+            x_offsets,
+            y_offsets,
+            z_offsets,
+        )
+
+        # pyre-fixme[6]: For 1st argument expected `Sequence[Union[int, SymInt]]`
+        #  but got `Tensor`.
+        c = torch.rand((z_lengths.sum()), device=x.device, dtype=x.dtype)
+        return c
+
+    @staticmethod
+    # pyre-fixme
+    def backward(ctx, grad_output: torch.Tensor):
+        (
+            x,
+            y,
+            x_lengths,
+            y_lengths,
+            z_lengths,
+            x_offsets,
+            y_offsets,
+            z_offsets,
+        ) = ctx.saved_tensors
+
+        grad_x = torch.rand(x.size(), device=x.device, dtype=x.dtype)
+        grad_y = torch.rand(y.size(), device=y.device, dtype=y.dtype)
+        return grad_x, grad_y, None, None, None, None, None, None, None, None
+
+
+# pyre-fixme[3]: Return type must be annotated.
+def meta_jagged_jagged_bmm_jagged_out(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    x_lengths: torch.Tensor,
+    x_offsets: torch.Tensor,
+    y_lengths: torch.Tensor,
+    y_offsets: torch.Tensor,
+    z_lengths: torch.Tensor,
+    z_offsets: torch.Tensor,
+    max_seq_len: int,
+    allow_tf32: bool = True,
+):
+    return JaggedJaggedBmmNoPaddingMeta.apply(
+        x,
+        y,
+        x_lengths,
+        x_offsets,
+        y_lengths,
+        y_offsets,
+        z_lengths,
+        z_offsets,
+        max_seq_len,
+        allow_tf32,
+    )

fbgemm_gpu/sll/triton/__init__.py

@@ -0,0 +1,127 @@
+#!/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.triton_dense_jagged_cat_jagged_out import (
+    dense_jagged_cat_jagged_out,
+)
+
+from fbgemm_gpu.sll.triton.triton_jagged2_to_padded_dense import (  # noqa F401
+    jagged2_to_padded_dense,
+    Jagged2ToPaddedDense,  # noqa F401
+)
+
+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_bmm_jagged_out import (  # noqa F401
+    array_jagged_bmm_jagged_out,
+    ArrayJaggedBmmNopadding,  # noqa F401
+    jagged_jagged_bmm_jagged_out,
+    JaggedJaggedBmmNoPadding,  # noqa F401
+    triton_array_jagged_bmm_jagged_out,  # noqa F401
+    triton_jagged_jagged_bmm_jagged_out,  # 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_elementwise_mul_jagged_out import (  # noqa F401
+    jagged_dense_elementwise_mul_jagged_out,
+    JaggedDenseElementwiseMul,  # 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_self_substraction_jagged_out import (
+    triton_jagged_self_substraction_jagged_out,
+)
+
+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_dense_jagged_cat_jagged_out": {
+        "CUDA": dense_jagged_cat_jagged_out,
+    },
+    "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_jagged2_to_padded_dense": {
+        "CUDA": jagged2_to_padded_dense,
+        "AutogradCUDA": jagged2_to_padded_dense,
+    },
+    "sll_array_jagged_bmm_jagged_out": {
+        "CUDA": array_jagged_bmm_jagged_out,
+        "AutogradCUDA": array_jagged_bmm_jagged_out,
+    },
+    "sll_jagged_jagged_bmm_jagged_out": {
+        "CUDA": jagged_jagged_bmm_jagged_out,
+        "AutogradCUDA": jagged_jagged_bmm_jagged_out,
+    },
+    "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,
+    },
+    "sll_jagged_dense_flash_attention": {
+        "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,
+    },
+    "sll_jagged_self_substraction_jagged_out": {
+        "CUDA": triton_jagged_self_substraction_jagged_out,
+    },
+    "sll_jagged_dense_elementwise_mul_jagged_out": {
+        "CUDA": jagged_dense_elementwise_mul_jagged_out,
+        "AutogradCUDA": jagged_dense_elementwise_mul_jagged_out,
+    },
+}

fbgemm_gpu/sll/triton/common.py

@@ -0,0 +1,38 @@
+# 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 next_power_of_two(N: int) -> int:
+    if N > 4096:
+        raise Exception(f"{N} is too large that is not supported yet")
+
+    if N > 2048:
+        return 4096
+    elif N > 1024:
+        return 2048
+    elif N > 512:
+        return 1024
+    elif N > 256:
+        return 512
+    elif N > 128:
+        return 256
+    elif N > 64:
+        return 128
+    elif N > 32:
+        return 64
+    else:
+        return 32
+
+
+def expect_contiguous(x: torch.Tensor) -> torch.Tensor:
+    if not x.is_contiguous():
+        return x.contiguous()
+    else:
+        return x

fbgemm_gpu/sll/triton/triton_dense_jagged_cat_jagged_out.py

@@ -0,0 +1,72 @@
+# 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
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def dense_jagged_cat_jagged_out_kernel(
+    a_ptr,  # dense
+    b_ptr,  # jagged
+    c_ptr,  # jagged
+    b_offsets_ptr,
+    c_offsets_ptr,
+    max_seq_len,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid_batch = tl.program_id(0)
+    b_start = tl.load(b_offsets_ptr + pid_batch)
+    b_end = tl.load(b_offsets_ptr + pid_batch + 1)
+    c_start = b_start + pid_batch
+    N = b_end - b_start
+    N = tl.minimum(N, max_seq_len)
+
+    a = tl.load(a_ptr + pid_batch)
+    tl.store(c_ptr + c_start, a)
+
+    offs_k = tl.arange(0, BLOCK_SIZE)
+    for k in range(0, N, BLOCK_SIZE):
+        b_offset = k + offs_k
+        b_ptrs = b_ptr + b_start + b_offset
+        b = tl.load(b_ptrs, mask=b_offset < N, other=0.0)
+        tl.store(c_ptr + c_start + 1 + b_offset, b, mask=b_offset < N)
+    tl.store(c_offsets_ptr + pid_batch, b_start + pid_batch)
+
+
+def dense_jagged_cat_jagged_out(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    b_offsets: torch.Tensor,
+    max_seq_len: int,
+):
+    assert a.is_contiguous()
+    assert b.is_contiguous()
+    assert b_offsets.is_contiguous()
+    B = a.size(0)
+    BLOCK_SIZE = 128
+    c = torch.zeros(b.size(0) + a.size(0), dtype=a.dtype, device=a.device)
+    c_offsets = torch.empty(
+        b_offsets.size(0), dtype=b_offsets.dtype, device=b_offsets.device
+    )  # B + 1
+
+    dense_jagged_cat_jagged_out_kernel[(B,)](
+        a,
+        b,
+        c,
+        b_offsets,
+        c_offsets,
+        max_seq_len,
+        # pyre-fixme[6]: For 7th argument expected `constexpr` but got `int`.
+        BLOCK_SIZE,
+    )
+
+    c_offsets[-1] = b_offsets[-1] + B
+
+    return c, c_offsets