blksprs 1.11__py3-none-any.whl → 2.0__py3-none-any.whl

blksprs/ops/partitioning.py

@@ -1,15 +1,16 @@
 import torch
 from torch import Tensor
+from torch._library import triton_op
 
-from blksprs.ops.flow import flow_forward_pull
+from blksprs.ops.flow import flow_pull_forward
 from blksprs.utils.blksprs_tensor import BlksprsTensor
-
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
-    validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
+    validate_sparsity, validate_sparsity_block_size
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
-          dim: int, sparsity_block_size: int, triton_block_size: int = None, lut: dict = None) -> (
+          dim: int, sparsity_block_size: int, lut: dict = None) -> (
         BlksprsTensor, Tensor):
     """Splits a block-sparse tensor in compressed form along the last dimension into partitions.
 
@@ -19,7 +20,6 @@ def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
         partitions (int): The number of partitions to split the block-sparse tensor into.
         dim (int): The dimension along which to split the tensor. Currently only supports dim=2.
         sparsity_block_size (int): The size of the sparsity blocks.
-        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
         lut (dict, optional): A dictionary containing the look-up tables for the operation (default ``None``).
 
     Returns:
@@ -34,83 +34,88 @@ def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
     validate_device(x)
     validate_sparsity(sparsity_block_size, (x, sparsity_layout))
     validate_sparsity_block_size(sparsity_block_size, x)
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
 
     adjusted_dim = dim % 3
     if adjusted_dim != 2:
         raise NotImplementedError("Currently only supports dim=2")
 
-    lut = _BlocksparseSplit.build_lut(lut, sparsity_layout, partitions)
+    lut = split_build_lut(lut, sparsity_layout, partitions)
+
+    return BlksprsTensor(split_forward(
+        x, lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"],
+        partitions, adjusted_dim, sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_output"]
+
+
+@triton_op("blksprs::split_forward", mutates_args={})
+def split_forward(x: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
+                  _: int, __: int, sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
+    with torch.no_grad():
+        return flow_pull_forward(x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
+                                 n_sparse_blocks)
 
-    return BlksprsTensor(
-        _BlocksparseSplit.apply(x, lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"],
-                                partitions, adjusted_dim, sparsity_block_size, lut["n_sparse_blocks"],
-                                triton_block_size)), lut["sparsity_layout_output"]
 
+def split_wrapper_backward(ctx, grad_output):
+    sparsity_layout = ctx.saved_tensors[0]
+    num_partitions = ctx.num_partitions
+    dim = ctx.dim
+    sparsity_block_size = ctx.sparsity_block_size
 
-class _BlocksparseSplit(torch.autograd.Function):
+    return merge(grad_output, sparsity_layout, num_partitions, dim,
+                 sparsity_block_size)[0], None, None, None, None, None, None, None
 
-    @staticmethod
-    def build_lut(lut: dict, sparsity_layout: Tensor, partitions: int):
-        if lut is None:
-            lut = dict()
 
-        if "sparsity_layout_output" not in lut:
-            sparsity_layout_output = (sparsity_layout
-                                      .reshape(sparsity_layout.size(0), sparsity_layout.size(1), partitions,
-                                               sparsity_layout.size(2) // partitions)
-                                      .permute(0, 2, 1, 3)
-                                      .reshape(sparsity_layout.size(0) * partitions, sparsity_layout.size(1),
-                                               sparsity_layout.size(2) // partitions).contiguous())
-            lut["sparsity_layout_output"] = sparsity_layout_output
+def split_build_lut(lut: dict, sparsity_layout: Tensor, partitions: int):
+    if lut is None:
+        lut = dict()
 
-        if "sparsity_lut" not in lut:
-            sparsity_lut = torch.nonzero(lut["sparsity_layout_output"]).contiguous()
-            lut["sparsity_lut"] = sparsity_lut
+    if "sparsity_layout_output" not in lut:
+        sparsity_layout_output = (sparsity_layout
+                                  .reshape(sparsity_layout.size(0), sparsity_layout.size(1), partitions,
+                                           sparsity_layout.size(2) // partitions)
+                                  .permute(0, 2, 1, 3)
+                                  .reshape(sparsity_layout.size(0) * partitions, sparsity_layout.size(1),
+                                           sparsity_layout.size(2) // partitions).contiguous())
+        lut["sparsity_layout_output"] = sparsity_layout_output
 
-        if "sparsity_reverse_lut" not in lut:
-            sparsity_layout_flat = sparsity_layout.reshape(-1)
-            sparsity_reverse_lut = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
-                                     (sparsity_layout_flat == 1) -
-                                     (1 * (sparsity_layout_flat == 0)))
-                                    .reshape(sparsity_layout.size(0), sparsity_layout.size(1), partitions,
-                                             sparsity_layout.size(2) // partitions)
-                                    .permute(0, 2, 1, 3).reshape(-1).contiguous())
-            lut["sparsity_reverse_lut"] = sparsity_reverse_lut
+    if "sparsity_lut" not in lut:
+        sparsity_lut = torch.nonzero(lut["sparsity_layout_output"]).contiguous()
+        lut["sparsity_lut"] = sparsity_lut
 
-        if "n_sparse_blocks" not in lut:
-            n_sparse_blocks = torch.sum(lut["sparsity_layout_output"].to(torch.int)).item()
-            lut["n_sparse_blocks"] = n_sparse_blocks
+    if "sparsity_reverse_lut" not in lut:
+        sparsity_layout_flat = sparsity_layout.reshape(-1)
+        sparsity_reverse_lut = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
+                                 (sparsity_layout_flat == 1) -
+                                 (1 * (sparsity_layout_flat == 0)))
+                                .reshape(sparsity_layout.size(0), sparsity_layout.size(1), partitions,
+                                         sparsity_layout.size(2) // partitions)
+                                .permute(0, 2, 1, 3).reshape(-1).contiguous())
+        lut["sparsity_reverse_lut"] = sparsity_reverse_lut
 
-        validate_contiguous(lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
+    if "n_sparse_blocks" not in lut:
+        n_sparse_blocks = torch.sum(lut["sparsity_layout_output"].to(torch.int)).item()
+        lut["n_sparse_blocks"] = n_sparse_blocks
 
-        return lut
+    validate_contiguous(lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
 
-    @staticmethod
-    def forward(ctx, x: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
-                num_partitions: int, dim: int, sparsity_block_size: int, n_sparse_blocks: int,
-                triton_block_size: int) -> Tensor:
-        ctx.save_for_backward(sparsity_layout_o)
-        ctx.num_partitions = num_partitions
-        ctx.dim = dim
+    return lut
 
-        return flow_forward_pull(ctx, x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
-                                 n_sparse_blocks, triton_block_size)
 
-    @staticmethod
-    def backward(ctx, grad_output):
-        sparsity_layout = ctx.saved_tensors[0]
-        num_partitions = ctx.num_partitions
-        dim = ctx.dim
-        sparsity_block_size = ctx.sparsity_block_size
-        triton_block_size = ctx.triton_block_size
+# noinspection PyUnusedLocal
+def split_setup_context(ctx, inputs, output):
+    (_, sparsity_layout_o, _, _, num_partitions, dim, sparsity_block_size, _) = inputs
 
-        return merge(grad_output, sparsity_layout, num_partitions, dim,
-                     sparsity_block_size, triton_block_size)[0], None, None, None, None, None, None, None, None
+    ctx.save_for_backward(sparsity_layout_o)
+    ctx.num_partitions = num_partitions
+    ctx.dim = dim
+    ctx.sparsity_block_size = sparsity_block_size
 
 
+split_forward.register_autograd(split_wrapper_backward, setup_context=split_setup_context)
+
+
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
-          dim: int, sparsity_block_size: int, triton_block_size: int = None, lut: dict = None) -> (
+          dim: int, sparsity_block_size: int, lut: dict = None) -> (
         BlksprsTensor, Tensor):
     """Merges the specified partitions of a block-sparse tensor in compressed form along the last dimension.
 
@@ -120,7 +125,6 @@ def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
         partitions (int): The number of partitions to be merged.
         dim (int): The dimension along which to merge the tensor. Currently only supports dim=2.
         sparsity_block_size (int): The size of the sparsity blocks.
-        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
         lut (dict, optional): A dictionary containing the look-up tables for the operation (default ``None``).
 
     Returns:
@@ -135,79 +139,83 @@ def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
     validate_device(x)
     validate_sparsity(sparsity_block_size, (x, sparsity_layout))
     validate_sparsity_block_size(sparsity_block_size, x)
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
 
     adjusted_dim = dim % 3
     if adjusted_dim != 2:
         raise NotImplementedError("Currently only supports dim=2")
 
-    lut = _BlocksparseMerge.build_lut(lut, sparsity_layout, partitions)
-
-    return BlksprsTensor(
-        _BlocksparseMerge.apply(x, lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"],
-                                partitions, adjusted_dim, sparsity_block_size, lut["n_sparse_blocks"],
-                                triton_block_size)), lut["sparsity_layout_output"]
-
-
-class _BlocksparseMerge(torch.autograd.Function):
-
-    @staticmethod
-    def build_lut(lut: dict, sparsity_layout: Tensor, partitions: int):
-        if lut is None:
-            lut = dict()
-
-        if "sparsity_layout_output" not in lut:
-            sparsity_layout_output = (sparsity_layout.reshape(sparsity_layout.size(0) // partitions, partitions,
-                                                              sparsity_layout.size(1), sparsity_layout.size(2))
-                                      .permute(0, 2, 1, 3)
-                                      .reshape(sparsity_layout.size(0) // partitions,
-                                               sparsity_layout.size(1),
-                                               sparsity_layout.size(2) * partitions).contiguous())
-            lut["sparsity_layout_output"] = sparsity_layout_output
-
-        if "sparsity_lut" not in lut:
-            sparsity_lut = torch.nonzero(lut["sparsity_layout_output"]).contiguous()
-            lut["sparsity_lut"] = sparsity_lut
-
-        if "sparsity_reverse_lut" not in lut:
-            sparsity_layout_flat = sparsity_layout.reshape(-1)
-            sparsity_reverse_lut = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
-                                     (sparsity_layout_flat == 1) -
-                                     (1 * (sparsity_layout_flat == 0)))
-                                    .reshape(sparsity_layout.size(0) // partitions, partitions,
-                                             sparsity_layout.size(1), sparsity_layout.size(2))
-                                    .permute(0, 2, 1, 3)
-                                    .reshape(sparsity_layout.size(0) // partitions,
-                                             sparsity_layout.size(1), sparsity_layout.size(2) * partitions)
-                                    .reshape(-1).contiguous())
-            lut["sparsity_reverse_lut"] = sparsity_reverse_lut
-
-        if "n_sparse_blocks" not in lut:
-            n_sparse_blocks = torch.sum(lut["sparsity_layout_output"].to(torch.int)).item()
-            lut["n_sparse_blocks"] = n_sparse_blocks
-
-        validate_contiguous(lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
-
-        return lut
-
-    @staticmethod
-    def forward(ctx, x: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
-                num_partitions: int, dim: int, sparsity_block_size: int, n_sparse_blocks: int,
-                triton_block_size: int) -> Tensor:
-        ctx.save_for_backward(sparsity_layout_o)
-        ctx.num_partitions = num_partitions
-        ctx.dim = dim
-
-        return flow_forward_pull(ctx, x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
-                                 n_sparse_blocks, triton_block_size)
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        sparsity_layout = ctx.saved_tensors[0]
-        num_partitions = ctx.num_partitions
-        dim = ctx.dim
-        sparsity_block_size = ctx.sparsity_block_size
-        triton_block_size = ctx.triton_block_size
-
-        return split(grad_output, sparsity_layout, num_partitions, dim,
-                     sparsity_block_size, triton_block_size)[0], None, None, None, None, None, None, None, None
+    lut = merge_build_lut(lut, sparsity_layout, partitions)
+
+    return BlksprsTensor(merge_forward(
+        x, lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"],
+        partitions, adjusted_dim, sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_output"]
+
+
+@triton_op("blksprs::merge_forward", mutates_args={})
+def merge_forward(x: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
+                  _: int, __: int, sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
+    with torch.no_grad():
+        return flow_pull_forward(x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
+                                 n_sparse_blocks)
+
+
+def merge_wrapper_backward(ctx, grad_output):
+    sparsity_layout = ctx.saved_tensors[0]
+    num_partitions = ctx.num_partitions
+    dim = ctx.dim
+    sparsity_block_size = ctx.sparsity_block_size
+
+    return split(grad_output, sparsity_layout, num_partitions, dim,
+                 sparsity_block_size)[0], None, None, None, None, None, None, None
+
+
+def merge_build_lut(lut: dict, sparsity_layout: Tensor, partitions: int):
+    if lut is None:
+        lut = dict()
+
+    if "sparsity_layout_output" not in lut:
+        sparsity_layout_output = (sparsity_layout.reshape(sparsity_layout.size(0) // partitions, partitions,
+                                                          sparsity_layout.size(1), sparsity_layout.size(2))
+                                  .permute(0, 2, 1, 3)
+                                  .reshape(sparsity_layout.size(0) // partitions,
+                                           sparsity_layout.size(1),
+                                           sparsity_layout.size(2) * partitions).contiguous())
+        lut["sparsity_layout_output"] = sparsity_layout_output
+
+    if "sparsity_lut" not in lut:
+        sparsity_lut = torch.nonzero(lut["sparsity_layout_output"]).contiguous()
+        lut["sparsity_lut"] = sparsity_lut
+
+    if "sparsity_reverse_lut" not in lut:
+        sparsity_layout_flat = sparsity_layout.reshape(-1)
+        sparsity_reverse_lut = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
+                                 (sparsity_layout_flat == 1) -
+                                 (1 * (sparsity_layout_flat == 0)))
+                                .reshape(sparsity_layout.size(0) // partitions, partitions,
+                                         sparsity_layout.size(1), sparsity_layout.size(2))
+                                .permute(0, 2, 1, 3)
+                                .reshape(sparsity_layout.size(0) // partitions,
+                                         sparsity_layout.size(1), sparsity_layout.size(2) * partitions)
+                                .reshape(-1).contiguous())
+        lut["sparsity_reverse_lut"] = sparsity_reverse_lut
+
+    if "n_sparse_blocks" not in lut:
+        n_sparse_blocks = torch.sum(lut["sparsity_layout_output"].to(torch.int)).item()
+        lut["n_sparse_blocks"] = n_sparse_blocks
+
+    validate_contiguous(lut["sparsity_layout_output"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
+
+    return lut
+
+
+# noinspection PyUnusedLocal
+def merge_setup_context(ctx, inputs, output):
+    (_, sparsity_layout_o, _, _, num_partitions, dim, sparsity_block_size, _) = inputs
+
+    ctx.save_for_backward(sparsity_layout_o)
+    ctx.num_partitions = num_partitions
+    ctx.dim = dim
+    ctx.sparsity_block_size = sparsity_block_size
+
+
+merge_forward.register_autograd(merge_wrapper_backward, setup_context=merge_setup_context)

blksprs/ops/repeat.py

@@ -1,17 +1,16 @@
 import torch
-import triton
 from torch import Tensor
+from torch._library import triton_op
 
-from blksprs.ops.flow import kernel_blocksparse_flow_push, flow_forward_pull, flow_forward_push
+from blksprs.ops.flow import flow_pull_forward, flow_push_forward
 from blksprs.utils.blksprs_tensor import BlksprsTensor
-from blksprs.utils.tools import get_triton_block_size, stride
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
-    validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
+    validate_sparsity, validate_sparsity_block_size
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int, int],
-           sparsity_block_size: int, sparsity_layout_output: Tensor = None, triton_block_size: int = None,
-           lut: dict = None) -> (
+           sparsity_block_size: int, sparsity_layout_output: Tensor = None, lut: dict = None) -> (
         BlksprsTensor, Tensor):
     """Repeats a block-spare tensor in compressed form according to the given repeats.
     
@@ -30,7 +29,6 @@ def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int,
             third dimension respectively.
         sparsity_block_size (int): The size of the sparsity blocks.
         sparsity_layout_output (Tensor): The desired sparsity layout of the output tensor (default ``None``).
-        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
         lut (dict, optional): A dictionary containing the look-up tables for the operation (default ``None``).
 
     Returns:
@@ -45,19 +43,17 @@ def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int,
     validate_device(x)
     validate_sparsity(sparsity_block_size, (x, sparsity_layout_x))
     validate_sparsity_block_size(sparsity_block_size, x)
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
 
-    lut = _BlocksparseRepeat.build_lut_repeat(lut, sparsity_layout_x, repeats, sparsity_layout_output)
+    lut = repeat_build_lut(lut, sparsity_layout_x, repeats, sparsity_layout_output)
 
-    return BlksprsTensor(
-        _BlocksparseRepeat.apply(x, sparsity_layout_x, lut["sparsity_layout_o"], lut["sparsity_lut"],
-                                 lut["sparsity_reverse_lut"], sparsity_block_size, lut["n_sparse_blocks"],
-                                 triton_block_size)), lut["sparsity_layout_o"]
+    return BlksprsTensor(repeat_forward(
+        x, sparsity_layout_x, lut["sparsity_layout_o"], lut["sparsity_lut"],
+        lut["sparsity_reverse_lut"], sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_o"]
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
-                      sparsity_block_size: int, sparsity_layout_output: Tensor = None,
-                      triton_block_size: int = None, lut: dict = None) -> (
+                      sparsity_block_size: int, sparsity_layout_output: Tensor = None, lut: dict = None) -> (
         BlksprsTensor, Tensor):
     """Repeats and interleaves the block-sparse tensor in compressed form.
 
@@ -74,7 +70,6 @@ def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
         repeats (int): The number of times to repeat the matrices.
         sparsity_block_size (int): The size of the sparsity blocks.
         sparsity_layout_output (Tensor): The desired sparsity layout of the output tensor (default ``None``).
-        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
         lut (dict, optional): A dictionary containing the look-up tables for the operation (default ``None``).
 
     Returns:
@@ -89,108 +84,111 @@ def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
     validate_device(x)
     validate_sparsity(sparsity_block_size, (x, sparsity_layout_x))
     validate_sparsity_block_size(sparsity_block_size, x)
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
 
-    lut = _BlocksparseRepeat.build_lut_repeat_interleave(lut, sparsity_layout_x, repeats, sparsity_layout_output)
-
-    return BlksprsTensor(
-        _BlocksparseRepeat.apply(x, sparsity_layout_x, lut["sparsity_layout_o"], lut["sparsity_lut"],
-                                 lut["sparsity_reverse_lut"], sparsity_block_size, lut["n_sparse_blocks"],
-                                 triton_block_size)), lut["sparsity_layout_o"]
-
-
-class _BlocksparseRepeat(torch.autograd.Function):
-
-    @staticmethod
-    def build_lut_repeat(lut: dict, sparsity_layout_x: Tensor, repeats: tuple[int, int, int],
-                         sparsity_layout_output: Tensor):
-        if lut is None:
-            lut = dict()
-
-        if "sparsity_layout_o" not in lut:
-            sparsity_layout_o = sparsity_layout_x.repeat(repeats[0], repeats[1], repeats[2])
-            lut["sparsity_layout_o"] = sparsity_layout_o
-
-        if sparsity_layout_output is not None:
-            sparsity_layout_o = torch.logical_and(lut["sparsity_layout_o"], sparsity_layout_output)
-            lut["sparsity_layout_o"] = sparsity_layout_o
-
-        if "sparsity_lut" not in lut:
-            sparsity_lut = torch.nonzero(lut["sparsity_layout_o"]).contiguous()
-            lut["sparsity_lut"] = sparsity_lut
-
-        if "sparsity_reverse_lut" not in lut:
-            sparsity_layout_flat = sparsity_layout_x.reshape(-1)
-            sparsity_reverse_lut = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
-                                     (sparsity_layout_flat == 1) -
-                                     (1 * (sparsity_layout_flat == 0)))
-                                    .reshape(sparsity_layout_x.size())
-                                    .repeat(repeats[0], repeats[1], repeats[2])
-                                    .reshape(-1).contiguous())
-            lut["sparsity_reverse_lut"] = sparsity_reverse_lut
-
-        if "n_sparse_blocks" not in lut:
-            n_sparse_blocks = torch.sum(lut["sparsity_layout_o"].to(torch.int)).item()
-            lut["n_sparse_blocks"] = n_sparse_blocks
-
-        validate_contiguous(sparsity_layout_o, lut["sparsity_lut"], lut["sparsity_reverse_lut"])
-
-        return lut
-
-    @staticmethod
-    def build_lut_repeat_interleave(lut: dict, sparsity_layout_x: Tensor, repeats: int,
-                                    sparsity_layout_output: Tensor):
-        if lut is None:
-            lut = dict()
-
-        if "sparsity_layout_o" not in lut:
-            sparsity_layout_o = torch.repeat_interleave(sparsity_layout_x, repeats, dim=0).contiguous()
-            lut["sparsity_layout_o"] = sparsity_layout_o
-
-        if sparsity_layout_output is not None:
-            sparsity_layout_o = torch.logical_and(lut["sparsity_layout_o"], sparsity_layout_output)
-            lut["sparsity_layout_o"] = sparsity_layout_o
-
-        if "sparsity_lut" not in lut:
-            sparsity_lut = torch.nonzero(lut["sparsity_layout_o"]).contiguous()
-            lut["sparsity_lut"] = sparsity_lut
-
-        if "sparsity_reverse_lut" not in lut:
-            sparsity_layout_flat = sparsity_layout_x.reshape(-1)
-            sparsity_reverse_lut = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
-                                     (sparsity_layout_flat == 1) -
-                                     (1 * (sparsity_layout_flat == 0)))
-                                    .reshape(sparsity_layout_x.size())
-                                    .repeat_interleave(repeats, dim=0)
-                                    .reshape(-1).contiguous())
-            lut["sparsity_reverse_lut"] = sparsity_reverse_lut
-
-        if "n_sparse_blocks" not in lut:
-            n_sparse_blocks = torch.sum(lut["sparsity_layout_o"].to(torch.int)).item()
-            lut["n_sparse_blocks"] = n_sparse_blocks
-
-        validate_contiguous(sparsity_layout_o, lut["sparsity_lut"], lut["sparsity_reverse_lut"])
-
-        return lut
-
-    @staticmethod
-    def forward(ctx, x: Tensor, sparsity_layout_x: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor,
-                sparsity_reverse_lut: Tensor,
-                sparsity_block_size: int, n_sparse_blocks: int,
-                triton_block_size: int) -> Tensor:
-        ctx.save_for_backward(sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut)
-
-        return flow_forward_pull(ctx, x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
-                                 n_sparse_blocks, triton_block_size)
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut = ctx.saved_tensors
-        sparsity_block_size = ctx.sparsity_block_size
-        triton_block_size = ctx.triton_block_size
-
-        n_sparse_blocks = torch.sum(sparsity_layout_x.to(torch.int)).item()
-
-        return flow_forward_push(None, grad_output, sparsity_layout_o, sparsity_lut,
-                                 sparsity_reverse_lut, sparsity_block_size, n_sparse_blocks,
-                                 triton_block_size), None, None, None, None, None, None, None
+    lut = repeat_interleave_build_lut(lut, sparsity_layout_x, repeats, sparsity_layout_output)
+
+    return BlksprsTensor(repeat_forward(
+        x, sparsity_layout_x, lut["sparsity_layout_o"], lut["sparsity_lut"],
+        lut["sparsity_reverse_lut"], sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_o"]
+
+
+@triton_op("blksprs::repeat_forward", mutates_args={})
+def repeat_forward(x: Tensor, _: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor,
+                   sparsity_reverse_lut: Tensor,
+                   sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
+    with torch.no_grad():
+        return flow_pull_forward(x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
+                                 n_sparse_blocks)
+
+
+def repeat_wrapper_backward(ctx, grad_output):
+    sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut = ctx.saved_tensors
+    sparsity_block_size = ctx.sparsity_block_size
+    n_sparse_blocks = torch.sum(sparsity_layout_x.to(torch.int)).item()
+
+    return flow_push_forward(grad_output, sparsity_layout_o, sparsity_lut,
+                             sparsity_reverse_lut, sparsity_block_size,
+                             n_sparse_blocks), None, None, None, None, None, None
+
+
+def repeat_build_lut(lut: dict, sparsity_layout_x: Tensor, repeats: tuple[int, int, int],
+                     sparsity_layout_output: Tensor):
+    if lut is None:
+        lut = dict()
+
+    if "sparsity_layout_o" not in lut:
+        sparsity_layout_o = sparsity_layout_x.repeat(repeats[0], repeats[1], repeats[2])
+        lut["sparsity_layout_o"] = sparsity_layout_o
+
+    if sparsity_layout_output is not None:
+        sparsity_layout_o = torch.logical_and(lut["sparsity_layout_o"], sparsity_layout_output)
+        lut["sparsity_layout_o"] = sparsity_layout_o
+
+    if "sparsity_lut" not in lut:
+        sparsity_lut = torch.nonzero(lut["sparsity_layout_o"]).contiguous()
+        lut["sparsity_lut"] = sparsity_lut
+
+    if "sparsity_reverse_lut" not in lut:
+        sparsity_layout_flat = sparsity_layout_x.reshape(-1)
+        sparsity_reverse_lut = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
+                                 (sparsity_layout_flat == 1) -
+                                 (1 * (sparsity_layout_flat == 0)))
+                                .reshape(sparsity_layout_x.size())
+                                .repeat(repeats[0], repeats[1], repeats[2])
+                                .reshape(-1).contiguous())
+        lut["sparsity_reverse_lut"] = sparsity_reverse_lut
+
+    if "n_sparse_blocks" not in lut:
+        n_sparse_blocks = torch.sum(lut["sparsity_layout_o"].to(torch.int)).item()
+        lut["n_sparse_blocks"] = n_sparse_blocks
+
+    validate_contiguous(sparsity_layout_o, lut["sparsity_lut"], lut["sparsity_reverse_lut"])
+
+    return lut
+
+
+def repeat_interleave_build_lut(lut: dict, sparsity_layout_x: Tensor, repeats: int,
+                                sparsity_layout_output: Tensor):
+    if lut is None:
+        lut = dict()
+
+    if "sparsity_layout_o" not in lut:
+        sparsity_layout_o = torch.repeat_interleave(sparsity_layout_x, repeats, dim=0).contiguous()
+        lut["sparsity_layout_o"] = sparsity_layout_o
+
+    if sparsity_layout_output is not None:
+        sparsity_layout_o = torch.logical_and(lut["sparsity_layout_o"], sparsity_layout_output)
+        lut["sparsity_layout_o"] = sparsity_layout_o
+
+    if "sparsity_lut" not in lut:
+        sparsity_lut = torch.nonzero(lut["sparsity_layout_o"]).contiguous()
+        lut["sparsity_lut"] = sparsity_lut
+
+    if "sparsity_reverse_lut" not in lut:
+        sparsity_layout_flat = sparsity_layout_x.reshape(-1)
+        sparsity_reverse_lut = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
+                                 (sparsity_layout_flat == 1) -
+                                 (1 * (sparsity_layout_flat == 0)))
+                                .reshape(sparsity_layout_x.size())
+                                .repeat_interleave(repeats, dim=0)
+                                .reshape(-1).contiguous())
+        lut["sparsity_reverse_lut"] = sparsity_reverse_lut
+
+    if "n_sparse_blocks" not in lut:
+        n_sparse_blocks = torch.sum(lut["sparsity_layout_o"].to(torch.int)).item()
+        lut["n_sparse_blocks"] = n_sparse_blocks
+
+    validate_contiguous(sparsity_layout_o, lut["sparsity_lut"], lut["sparsity_reverse_lut"])
+
+    return lut
+
+
+# noinspection PyUnusedLocal
+def repeat_setup_context(ctx, inputs, output):
+    (_, sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut, sparsity_block_size, _) = inputs
+
+    ctx.save_for_backward(sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut)
+    ctx.sparsity_block_size = sparsity_block_size
+
+
+repeat_forward.register_autograd(repeat_wrapper_backward, setup_context=repeat_setup_context)