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

blksprs/ops/partitioning.py

@@ -1,15 +1,17 @@
 import torch
 from torch import Tensor
+from torch._library import triton_op
 
-from blksprs.ops.flow import flow_forward
+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) -> (BlksprsTensor, Tensor):
+          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.
 
     Args:
@@ -18,7 +20,7 @@ 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:
         BlksprsTensor: The block-sparse tensor split into partitions in compressed form.
@@ -32,63 +34,89 @@ 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)
 
-    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())
+    adjusted_dim = dim % 3
+    if adjusted_dim != 2:
+        raise NotImplementedError("Currently only supports dim=2")
 
-    sparsity_lut = torch.nonzero(sparsity_layout_output).contiguous()
+    lut = split_build_lut(lut, sparsity_layout, partitions)
 
-    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())
+    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"]
 
-    n_sparse_blocks = torch.sum(sparsity_layout_output.to(torch.int)).item()
 
-    validate_contiguous(sparsity_layout_output, sparsity_lut, sparsity_reverse_lut)
+@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)
 
-    adjusted_dim = dim % 3
-    if adjusted_dim != 2:
-        raise NotImplementedError("Currently only supports dim=2")
 
-    return BlksprsTensor(_BlocksparseSplit.apply(x, sparsity_layout_output, sparsity_lut, sparsity_reverse_lut, partitions,
-                                   adjusted_dim, sparsity_block_size, n_sparse_blocks, triton_block_size)), 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
+
+    return merge(grad_output, sparsity_layout, num_partitions, dim,
+                 sparsity_block_size)[0], None, None, None, None, None, None, None
+
+
+def split_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
+
+    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), 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 "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"])
 
-class _BlocksparseSplit(torch.autograd.Function):
+    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(ctx, x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut, sparsity_block_size,
-                            n_sparse_blocks, triton_block_size)
+# noinspection PyUnusedLocal
+def split_setup_context(ctx, inputs, output):
+    (_, sparsity_layout_o, _, _, num_partitions, dim, sparsity_block_size, _) = inputs
 
-    @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
+    ctx.save_for_backward(sparsity_layout_o)
+    ctx.num_partitions = num_partitions
+    ctx.dim = dim
+    ctx.sparsity_block_size = sparsity_block_size
 
-        return merge(grad_output, sparsity_layout, num_partitions, dim,
-                     sparsity_block_size, triton_block_size)[0], None, None, None, None, None, None, None, None
 
+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) -> (BlksprsTensor, Tensor):
+          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.
 
     Args:
@@ -97,7 +125,7 @@ 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:
         BlksprsTensor: The merged block-sparse tensor in compressed form.
@@ -111,60 +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)
 
-    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())
+    adjusted_dim = dim % 3
+    if adjusted_dim != 2:
+        raise NotImplementedError("Currently only supports dim=2")
 
-    sparsity_lut = torch.nonzero(sparsity_layout_output).contiguous()
+    lut = merge_build_lut(lut, sparsity_layout, partitions)
 
-    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())
+    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"]
 
-    n_sparse_blocks = torch.sum(sparsity_layout_output.to(torch.int)).item()
 
-    validate_contiguous(sparsity_layout_output, sparsity_lut, sparsity_reverse_lut)
+@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)
 
-    adjusted_dim = dim % 3
-    if adjusted_dim != 2:
-        raise NotImplementedError("Currently only supports dim=2")
 
-    return BlksprsTensor(_BlocksparseMerge.apply(x, sparsity_layout_output, sparsity_lut, sparsity_reverse_lut, partitions,
-                                   adjusted_dim, sparsity_block_size, n_sparse_blocks, triton_block_size)), sparsity_layout_output
+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
 
-class _BlocksparseMerge(torch.autograd.Function):
+    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(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 merge_setup_context(ctx, inputs, output):
+    (_, sparsity_layout_o, _, _, num_partitions, dim, sparsity_block_size, _) = inputs
 
-        return split(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
 
 
+merge_forward.register_autograd(merge_wrapper_backward, setup_context=merge_setup_context)

blksprs/ops/repeat.py

@@ -1,16 +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
+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) -> (
+           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.
     
@@ -29,7 +29,7 @@ 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:
         BlksprsTensor: A block-sparse tensor in compressed form containing the repeated values.
@@ -43,35 +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)
 
-    sparsity_layout_o = sparsity_layout_x.repeat(repeats[0], repeats[1], repeats[2])
+    lut = repeat_build_lut(lut, sparsity_layout_x, repeats, sparsity_layout_output)
 
-    if sparsity_layout_output is not None:
-        sparsity_layout_o = torch.logical_and(sparsity_layout_o, sparsity_layout_output)
-
-    sparsity_lut = torch.nonzero(sparsity_layout_o).contiguous()
-
-    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())
-
-    n_sparse_blocks = torch.sum(sparsity_layout_o.to(torch.int)).item()
-
-    validate_contiguous(sparsity_layout_o, sparsity_lut, sparsity_reverse_lut)
-
-    return BlksprsTensor(
-        _BlocksparseRepeat.apply(x, sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
-                                 sparsity_block_size, n_sparse_blocks, triton_block_size)), 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) -> (
+                      sparsity_block_size: int, sparsity_layout_output: Tensor = None, lut: dict = None) -> (
         BlksprsTensor, Tensor):
     """Repeats and interleaves the block-sparse tensor in compressed form.
 
@@ -88,7 +70,7 @@ 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:
         BlksprsTensor: A block-sparse tensor in compressed form containing the repeated and interleaved matrices.
@@ -102,83 +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)
 
-    sparsity_layout_o = torch.repeat_interleave(sparsity_layout_x, repeats, dim=0).contiguous()
+    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(sparsity_layout_o, sparsity_layout_output)
-
-    sparsity_lut = torch.nonzero(sparsity_layout_o).contiguous()
-
-    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())
-
-    n_sparse_blocks = torch.sum(sparsity_layout_o.to(torch.int)).item()
-
-    validate_contiguous(sparsity_layout_o, sparsity_lut, sparsity_reverse_lut)
-
-    return BlksprsTensor(
-        _BlocksparseRepeat.apply(x, sparsity_layout_x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
-                                 sparsity_block_size, n_sparse_blocks, triton_block_size)), sparsity_layout_o
-
-
-class _BlocksparseRepeat(torch.autograd.Function):
-
-    @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)
-        ctx.x_size = x.size()
-        ctx.x_stride = stride(x)
-
-        return flow_forward(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
-        x_size = ctx.x_size
-        x_stride = ctx.x_stride
-        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()
-
-        output = torch.zeros(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
-                             dtype=grad_output.dtype, device=grad_output.device)
-
-        x_b, x_r, x_c = grad_output.size()
-        x_b_s, x_r_s, x_c_s = stride(grad_output)
-        s_l_x_b, s_l_x_r, s_l_x_c = sparsity_layout_o.size()
-        s_l_x_b_s, s_l_x_r_s, s_l_x_c_s = stride(sparsity_layout_o)
-        s_lut_r, s_lut_c = sparsity_lut.size()
-        s_lut_r_s, s_lut_c_s = stride(sparsity_lut)
-        o_b, o_r, o_c = x_size
-        o_b_s, o_r_s, o_c_s = x_stride
-
-        if triton_block_size is None:
-            triton_block_size = get_triton_block_size(sparsity_block_size)
-
-        triton_grid = lambda meta: [x_b,
-                                    triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
-                                    triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
-
-        (kernel_blocksparse_flow_push[triton_grid]
-         (grad_output,
-          x_b, x_b_s, x_r_s, x_c_s,
-          s_l_x_b, s_l_x_b_s, s_l_x_r_s, s_l_x_c_s,
-          sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
-          sparsity_reverse_lut,
-          output,
-          o_b, o_b_s, o_r_s, o_c_s,
-          triton_block_size))
-
-        return output, None, None, None, None, None, None, 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)