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

fbgemm_gpu/sll/triton/triton_jagged_dense_elementwise_mul_jagged_out.py

@@ -0,0 +1,175 @@
+# 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 jagged_dense_elementwise_mul_jagged_out_kernel(
+    a_ptr,  # 1d jagged
+    b_ptr,  # dense
+    c_ptr,  # 1d jagged
+    a_seq_lengths_ptr,
+    a_offsets_ptr,
+    stride_a,
+    stride_bm,
+    stride_bn,
+    max_seq_len,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_batch = tl.program_id(0)
+    pid_row_block = tl.program_id(1)
+
+    batch_offset = tl.load(a_offsets_ptr + pid_batch)
+    batch_seq_len = tl.load(a_seq_lengths_ptr + pid_batch)
+    truncated_seq_len = tl.minimum(batch_seq_len, max_seq_len)
+
+    offs_row = tl.arange(0, BLOCK_M)
+    offs_col = tl.arange(0, BLOCK_N)
+
+    rows = pid_row_block * BLOCK_M + offs_row
+
+    # a start + batch offset + row offsets + initial col offsets
+    a_ptrs = (
+        a_ptr
+        + batch_offset * stride_a
+        + rows[:, None] * truncated_seq_len
+        + offs_col[None, :]
+    )
+
+    # b start + row offsets + initial col offsets
+    b_ptrs = b_ptr + rows[:, None] * stride_bm + offs_col[None, :] * stride_bn
+
+    # c start + batch offset + row offsets + initial col offsets
+    c_ptrs = (
+        c_ptr + batch_offset + rows[:, None] * truncated_seq_len + offs_col[None, :]
+    )
+
+    for block_start in range(0, truncated_seq_len, BLOCK_N):
+        cols = block_start + offs_col
+        # pyre-fixme[16]: `int` has no attribute `__getitem__`.
+        mask = (rows[:, None] < truncated_seq_len) & (cols[None, :] < truncated_seq_len)
+        a = tl.load(a_ptrs, mask=mask)
+        a_ptrs += BLOCK_N
+
+        b = tl.load(b_ptrs, mask=mask)
+        b_ptrs += BLOCK_N
+
+        c = a * b
+        tl.store(c_ptrs, c, mask=mask)
+        c_ptrs += BLOCK_N
+
+
+def triton_jagged_dense_elementwise_mul_jagged_out(
+    jagged_A,
+    dense_B,
+    seq_lengths_a,
+    offsets_a,
+    max_seq_len,
+):
+    B = seq_lengths_a.size(0)
+    total_L = jagged_A.size(0)
+
+    jagged_C = torch.zeros((total_L), device=jagged_A.device, dtype=jagged_A.dtype)
+
+    BLOCK_M = 32
+    BLOCK_N = 32
+    num_blocks_m = triton.cdiv(max_seq_len, BLOCK_M)
+    grid = (B, num_blocks_m)
+
+    jagged_dense_elementwise_mul_jagged_out_kernel[grid](
+        jagged_A,
+        dense_B,
+        jagged_C,
+        seq_lengths_a,
+        offsets_a,
+        jagged_A.stride(0),
+        dense_B.stride(0),
+        dense_B.stride(1),
+        max_seq_len,
+        BLOCK_M,
+        BLOCK_N,
+    )
+
+    return jagged_C
+
+
+class JaggedDenseElementwiseMul(torch.autograd.Function):
+    """
+    Compute elementwise multiplication between jagged tensor and dense tensor.
+    z = x * y
+    x: [sum_B(L_i)]
+    y: dense tensor
+    z: [sum_B(L_i)]
+    """
+
+    @staticmethod
+    # pyre-fixme
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        x_seq_lengths: torch.Tensor,
+        x_offsets: torch.Tensor,
+        max_seq_len: int,
+    ):
+        ctx.max_seq_len = max_seq_len
+
+        ctx.save_for_backward(
+            x,
+            y,
+            x_seq_lengths,
+            x_offsets,
+        )
+
+        return triton_jagged_dense_elementwise_mul_jagged_out(
+            x,
+            y,
+            x_seq_lengths,
+            x_offsets,
+            max_seq_len,
+        )
+
+    @staticmethod
+    # pyre-fixme
+    def backward(ctx, grad_output: torch.Tensor):
+        (
+            x,
+            y,
+            x_seq_lengths,
+            x_offsets,
+        ) = ctx.saved_tensors
+
+        grad_x = triton_jagged_dense_elementwise_mul_jagged_out(
+            grad_output,
+            y,
+            x_seq_lengths,
+            x_offsets,
+            ctx.max_seq_len,
+        )
+
+        return grad_x, None, None, None, None
+
+
+def 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 JaggedDenseElementwiseMul.apply(
+        x,
+        y,
+        x_seq_lengths,
+        x_offsets,
+        max_seq_len,
+    )