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

blksprs/ops/partitioning.py

@@ -1,7 +1,7 @@
 import torch
 from torch import Tensor
 
-from blksprs.ops.flow import flow_forward
+from blksprs.ops.flow import flow_forward_pull
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
@@ -9,7 +9,8 @@ from blksprs.utils.validation import validate_dimensions, validate_contiguous, v
 
 
 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, triton_block_size: int = None, lut: dict = None) -> (
+        BlksprsTensor, Tensor):
     """Splits a block-sparse tensor in compressed form along the last dimension into partitions.
 
     Args:
@@ -19,6 +20,7 @@ def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
         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.
@@ -34,46 +36,66 @@ def split(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
     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())
-
-    sparsity_lut = torch.nonzero(sparsity_layout_output).contiguous()
-
-    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())
-
-    n_sparse_blocks = torch.sum(sparsity_layout_output.to(torch.int)).item()
-
-    validate_contiguous(sparsity_layout_output, sparsity_lut, sparsity_reverse_lut)
-
     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
+    lut = _BlocksparseSplit.build_lut(lut, sparsity_layout, partitions)
+
+    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"]
 
 
 class _BlocksparseSplit(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), 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"])
+
+        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:
+                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)
+        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):
@@ -88,7 +110,8 @@ class _BlocksparseSplit(torch.autograd.Function):
 
 
 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, triton_block_size: int = None, lut: dict = None) -> (
+        BlksprsTensor, Tensor):
     """Merges the specified partitions of a block-sparse tensor in compressed form along the last dimension.
 
     Args:
@@ -98,6 +121,7 @@ def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
         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.
@@ -113,48 +137,69 @@ def merge(x: BlksprsTensor, sparsity_layout: Tensor, partitions: int,
     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())
-
-    sparsity_lut = torch.nonzero(sparsity_layout_output).contiguous()
-
-    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())
-
-    n_sparse_blocks = torch.sum(sparsity_layout_output.to(torch.int)).item()
-
-    validate_contiguous(sparsity_layout_output, sparsity_lut, sparsity_reverse_lut)
-
     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
+    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:
+                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)
+        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):
@@ -166,5 +211,3 @@ class _BlocksparseMerge(torch.autograd.Function):
 
         return split(grad_output, sparsity_layout, num_partitions, dim,
                      sparsity_block_size, triton_block_size)[0], None, None, None, None, None, None, None, None
-
-

blksprs/ops/repeat.py

@@ -2,7 +2,7 @@ import torch
 import triton
 from torch import Tensor
 
-from blksprs.ops.flow import kernel_blocksparse_flow_push, flow_forward
+from blksprs.ops.flow import kernel_blocksparse_flow_push, flow_forward_pull, flow_forward_push
 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, \
@@ -10,7 +10,8 @@ from blksprs.utils.validation import validate_dimensions, validate_contiguous, v
 
 
 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, triton_block_size: int = None,
+           lut: dict = None) -> (
         BlksprsTensor, Tensor):
     """Repeats a block-spare tensor in compressed form according to the given repeats.
     
@@ -30,6 +31,7 @@ def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int,
         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.
@@ -45,33 +47,17 @@ def repeat(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: tuple[int, int,
     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])
-
-    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)
+    lut = _BlocksparseRepeat.build_lut_repeat(lut, sparsity_layout_x, repeats, sparsity_layout_output)
 
     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
+        _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"]
 
 
 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) -> (
+                      triton_block_size: int = None, lut: dict = None) -> (
         BlksprsTensor, Tensor):
     """Repeats and interleaves the block-sparse tensor in compressed form.
 
@@ -89,6 +75,7 @@ def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
         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.
@@ -104,31 +91,87 @@ def repeat_interleave(x: BlksprsTensor, sparsity_layout_x: Tensor, repeats: int,
     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 = _BlocksparseRepeat.build_lut_repeat_interleave(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)
+    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"]
 
-    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())
+class _BlocksparseRepeat(torch.autograd.Function):
 
-    n_sparse_blocks = torch.sum(sparsity_layout_o.to(torch.int)).item()
+    @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()
 
-    validate_contiguous(sparsity_layout_o, sparsity_lut, sparsity_reverse_lut)
+        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
 
-    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
+        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
 
-class _BlocksparseRepeat(torch.autograd.Function):
+        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,
@@ -136,49 +179,18 @@ class _BlocksparseRepeat(torch.autograd.Function):
                 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)
+        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
-        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
+        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