blksprs 2.0rc7__py3-none-any.whl → 2.0rc8__py3-none-any.whl

blksprs/__init__.py

@@ -1,3 +1,4 @@
+from blksprs.utils.tools import version
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
 

blksprs/layouting/distribution_layout.py

@@ -1,6 +1,10 @@
+import typing
+
 import torch
 import triton
 from torch import Tensor
+from torch._library import triton_op
+from torch._library.triton import wrap_triton
 from triton import language as tl
 
 from blksprs.utils.blksprs_tensor import BlksprsTensor
@@ -10,6 +14,7 @@ from blksprs.utils.validation import validate_dimensions, validate_device, \
     validate_contiguous
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: Tensor,
                               dim: int, size_target: torch.Size,
                               sparsity_block_size: int) -> Tensor:
@@ -34,32 +39,40 @@ def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: T
 
     adjusted_dim = dim % 3
 
-    output = torch.zeros(size_target[0], size_target[1] // sparsity_block_size, size_target[2] // sparsity_block_size,
-                         dtype=torch.bool, device=indices.device)
-
-    i_b, i_r, i_c = indices.size()
-    i_b_s, i_r_s, i_c_s = stride(indices)
-    s_lut_i_r, s_lut_i_c = sparsity_lut_i.size()
-    s_lut_i_r_s, s_lut_i_c_s = stride(sparsity_lut_i)
-    o_b, o_r, o_c = output.size()
-    o_b_s, o_r_s, o_c_s = stride(output)
-
-    triton_grid = lambda meta: [i_b,
-                                triton.cdiv(i_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(i_c, meta["TRITON_BLOCK_SIZE"])]
-
-    # TODO wrap
-    (build_distribution_layout_kernel[triton_grid]
-     (indices,
-      i_b, i_b_s, i_r_s, i_c_s,
-      sparsity_lut_i,
-      s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
-      adjusted_dim,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size))
-
-    return output
+    return build_distribution_layout_operation(indices, sparsity_lut_i, adjusted_dim, size_target, sparsity_block_size)
+
+
+@triton_op("blksprs::build_distribution_layout", mutates_args={})
+def build_distribution_layout_operation(indices: Tensor, sparsity_lut_i: Tensor,
+                                        adjusted_dim: int, size_target: typing.List[int],
+                                        sparsity_block_size: int) -> Tensor:
+    with torch.no_grad():
+        output = torch.zeros(size_target[0], size_target[1] // sparsity_block_size,
+                             size_target[2] // sparsity_block_size,
+                             dtype=torch.bool, device=indices.device)
+
+        i_b, i_r, i_c = indices.size()
+        i_b_s, i_r_s, i_c_s = stride(indices)
+        s_lut_i_r, s_lut_i_c = sparsity_lut_i.size()
+        s_lut_i_r_s, s_lut_i_c_s = stride(sparsity_lut_i)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
+
+        triton_grid = lambda meta: [i_b,
+                                    triton.cdiv(i_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(i_c, meta["TRITON_BLOCK_SIZE"])]
+
+        (wrap_triton(build_distribution_layout_kernel)[triton_grid]
+         (indices,
+          i_b, i_b_s, i_r_s, i_c_s,
+          sparsity_lut_i,
+          s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
+          adjusted_dim,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_block_size))
+
+        return output
 
 
 @triton.autotune(

blksprs/layouting/sparsity_layout.py

@@ -3,7 +3,7 @@ import math
 import torch
 import triton
 from torch import Tensor
-from torch._library.triton import wrap_triton
+from torch._library.triton import wrap_triton, triton_op
 from triton import language as tl
 
 from blksprs.utils.blksprs_tensor import BlksprsTensor
@@ -29,27 +29,32 @@ def build_sparsity_layout(x: Tensor, sparsity_block_size: int) -> Tensor:
     validate_contiguous(x)
     validate_device(x)
 
-    output = torch.zeros(x.size(0), x.size(1) // sparsity_block_size, x.size(2) // sparsity_block_size,
-                         dtype=torch.bool, device=x.device)
+    return build_sparsity_layout_operation(x, sparsity_block_size)
 
-    x_b, x_r, x_c = x.size()
-    x_b_s, x_r_s, x_c_s = stride(x)
-    o_b, o_r, o_c = output.size()
-    o_b_s, o_r_s, o_c_s = stride(output)
 
-    triton_grid = lambda meta: [x_b,
-                                triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+@triton_op("blksprs::build_sparsity_layout", mutates_args={})
+def build_sparsity_layout_operation(x: Tensor, sparsity_block_size: int) -> Tensor:
+    with torch.no_grad():
+        output = torch.zeros(x.size(0), x.size(1) // sparsity_block_size, x.size(2) // sparsity_block_size,
+                             dtype=torch.bool, device=x.device)
 
-    # TODO wrap
-    (build_sparsity_layout_kernel[triton_grid]
-     (x,
-      x_b, x_b_s, x_r_s, x_c_s,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size))
+        x_b, x_r, x_c = x.size()
+        x_b_s, x_r_s, x_c_s = stride(x)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
 
-    return output
+        triton_grid = lambda meta: [x_b,
+                                    triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+
+        (wrap_triton(build_sparsity_layout_kernel)[triton_grid]
+         (x,
+          x_b, x_b_s, x_r_s, x_c_s,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_block_size))
+
+        return output
 
 
 @triton.autotune(
@@ -87,6 +92,7 @@ def build_sparsity_layout_kernel(x,
         tl.store(o + blk_o_idx, 1, mask=blk_o_msk)
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def build_sparsity_layout_adaption(x: BlksprsTensor, sparsity_layout_from: Tensor,
                                    sparsity_block_size_from: int, sparsity_block_size_to: int) -> Tensor:
     """Builds the sparsity layout of a block-sparse tensor in compressed form if a different sparsity block size were
@@ -114,33 +120,40 @@ def build_sparsity_layout_adaption(x: BlksprsTensor, sparsity_layout_from: Tenso
 
     validate_contiguous(sparsity_layout_from, sparsity_lut)
 
-    o_b = sparsity_layout_from.size(0)
-    o_r = math.ceil(sparsity_layout_from.size(1) * sparsity_block_size_from // sparsity_block_size_to)
-    o_c = math.ceil(sparsity_layout_from.size(2) * sparsity_block_size_from // sparsity_block_size_to)
-
-    output = torch.zeros(o_b, o_r, o_c, dtype=torch.bool, device=x.device)
-
-    x_b, x_r, x_c = x.size()
-    x_b_s, x_r_s, x_c_s = stride(x)
-    s_lut_r, s_lut_c = sparsity_lut.size()
-    s_lut_r_s, s_lut_c_s = stride(sparsity_lut)
-    o_b_s, o_r_s, o_c_s = stride(output)
-
-    triton_grid = lambda meta: [x_b,
-                                triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
-
-    # TODO wrap
-    (build_sparsity_layout_adaption_kernel[triton_grid]
-     (x,
-      x_b, x_b_s, x_r_s, x_c_s,
-      sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size_from,
-      sparsity_block_size_to))
-
-    return output
+    return build_sparsity_layout_adaption_operation(x, sparsity_layout_from, sparsity_lut,
+                                                    sparsity_block_size_from, sparsity_block_size_to)
+
+
+@triton_op("blksprs::build_sparsity_layout_adaption", mutates_args={})
+def build_sparsity_layout_adaption_operation(x: Tensor, sparsity_layout_from: Tensor, sparsity_lut: Tensor,
+                                             sparsity_block_size_from: int, sparsity_block_size_to: int) -> Tensor:
+    with torch.no_grad():
+        o_b = sparsity_layout_from.size(0)
+        o_r = math.ceil(sparsity_layout_from.size(1) * sparsity_block_size_from // sparsity_block_size_to)
+        o_c = math.ceil(sparsity_layout_from.size(2) * sparsity_block_size_from // sparsity_block_size_to)
+
+        output = torch.zeros(o_b, o_r, o_c, dtype=torch.bool, device=x.device)
+
+        x_b, x_r, x_c = x.size()
+        x_b_s, x_r_s, x_c_s = stride(x)
+        s_lut_r, s_lut_c = sparsity_lut.size()
+        s_lut_r_s, s_lut_c_s = stride(sparsity_lut)
+        o_b_s, o_r_s, o_c_s = stride(output)
+
+        triton_grid = lambda meta: [x_b,
+                                    triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+
+        (wrap_triton(build_sparsity_layout_adaption_kernel)[triton_grid]
+         (x,
+          x_b, x_b_s, x_r_s, x_c_s,
+          sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_block_size_from,
+          sparsity_block_size_to))
+
+        return output
 
 
 @triton.autotune(

blksprs/ops/conversion.py

@@ -52,33 +52,34 @@ def to_sparse(x: Tensor, sparsity_layout: Tensor,
                                            lut["sparsity_lut"], sparsity_block_size, lut["n_sparse_blocks"]))
 
 
-@triton_op("blksprs::to_sparse", mutates_args={})
+@triton_op("blksprs::to_sparse_forward", mutates_args={})
 def to_sparse_forward(x: Tensor, _: Tensor,
                       sparsity_lut: Tensor, sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
-    output = torch.zeros(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
-                         dtype=x.dtype, device=x.device)
+    with torch.no_grad():
+        output = torch.zeros(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
+                             dtype=x.dtype, device=x.device)
 
-    x_b, x_r, x_c = x.size()
-    x_b_s, x_r_s, x_c_s = stride(x)
-    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 = output.size()
-    o_b_s, o_r_s, o_c_s = stride(output)
+        x_b, x_r, x_c = x.size()
+        x_b_s, x_r_s, x_c_s = stride(x)
+        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 = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
 
-    triton_grid = lambda meta: [o_b,
-                                triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
+        triton_grid = lambda meta: [o_b,
+                                    triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
 
-    (wrap_triton(to_sparse_kernel)[triton_grid]
-     (x, x_b, x_b_s, x_r_s, x_c_s,
-      sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
-      output, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size))
+        (wrap_triton(to_sparse_kernel)[triton_grid]
+         (x, x_b, x_b_s, x_r_s, x_c_s,
+          sparsity_lut, s_lut_r, s_lut_r_s, s_lut_c_s,
+          output, o_b_s, o_r_s, o_c_s,
+          sparsity_block_size))
 
-    return output
+        return output
 
 
-def to_sparse_backward(ctx, grad_output):
+def to_sparse_wrapper_backward(ctx, grad_output):
     sparsity_layout = ctx.saved_tensors[0]
     sparsity_block_size = ctx.sparsity_block_size
 
@@ -161,7 +162,7 @@ def to_sparse_setup_context(ctx, inputs, output):
     ctx.sparsity_block_size = sparsity_block_size
 
 
-to_sparse_forward.register_autograd(to_sparse_backward, setup_context=to_sparse_setup_context)
+to_sparse_forward.register_autograd(to_sparse_wrapper_backward, setup_context=to_sparse_setup_context)
 
 
 def from_blksprs(x: BlksprsTensor, sparsity_layout: Tensor,
@@ -207,38 +208,39 @@ def to_dense(x: BlksprsTensor, sparsity_layout: Tensor,
                             lut["sparsity_reverse_lut"], sparsity_block_size, fill_value)
 
 
-@triton_op("blksprs::to_dense", mutates_args={})
+@triton_op("blksprs::to_dense_forward", mutates_args={})
 def to_dense_forward(x: Tensor, sparsity_layout: Tensor,
                      sparsity_reverse_lut: Tensor,
                      sparsity_block_size: int, fill_value: float) -> Tensor:
-    output = torch.full(size=(sparsity_layout.size(0), sparsity_layout.size(1) * sparsity_block_size,
-                              sparsity_layout.size(2) * sparsity_block_size), fill_value=fill_value,
-                        dtype=x.dtype, device=x.device)
-
-    x_b, x_r, x_c = x.shape
-    x_b_s, x_r_s, x_c_s = stride(x)
-    s_l_b, s_l_r, s_l_c = sparsity_layout.size()
-    s_l_b_s, s_l_r_s, s_l_c_s = stride(sparsity_layout)
-    o_b, o_r, o_c = output.size()
-    o_b_s, o_r_s, o_c_s = stride(output)
-
-    triton_grid = lambda meta: [o_b,
-                                triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
-
-    (wrap_triton(to_dense_kernel)[triton_grid]
-     (x,
-      x_b, x_b_s, x_r_s, x_c_s,
-      s_l_b, s_l_b_s, s_l_r_s, s_l_c_s,
-      sparsity_reverse_lut,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size))
-
-    return output
-
-
-def to_dense_backward(ctx, grad_output):
+    with torch.no_grad():
+        output = torch.full(size=(sparsity_layout.size(0), sparsity_layout.size(1) * sparsity_block_size,
+                                  sparsity_layout.size(2) * sparsity_block_size), fill_value=fill_value,
+                            dtype=x.dtype, device=x.device)
+
+        x_b, x_r, x_c = x.shape
+        x_b_s, x_r_s, x_c_s = stride(x)
+        s_l_b, s_l_r, s_l_c = sparsity_layout.size()
+        s_l_b_s, s_l_r_s, s_l_c_s = stride(sparsity_layout)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
+
+        triton_grid = lambda meta: [o_b,
+                                    triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
+
+        (wrap_triton(to_dense_kernel)[triton_grid]
+         (x,
+          x_b, x_b_s, x_r_s, x_c_s,
+          s_l_b, s_l_b_s, s_l_r_s, s_l_c_s,
+          sparsity_reverse_lut,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_block_size))
+
+        return output
+
+
+def to_dense_wrapper_backward(ctx, grad_output):
     sparsity_layout = ctx.saved_tensors[0]
     sparsity_block_size = ctx.sparsity_block_size
 
@@ -316,7 +318,7 @@ def to_dense_setup_context(ctx, inputs, output):
     ctx.sparsity_block_size = sparsity_block_size
 
 
-to_dense_forward.register_autograd(to_dense_backward, setup_context=to_dense_setup_context)
+to_dense_forward.register_autograd(to_dense_wrapper_backward, setup_context=to_dense_setup_context)
 
 
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
@@ -372,45 +374,45 @@ def adapt_layout(x: BlksprsTensor, sparsity_layout_from: Tensor, sparsity_block_
                                               n_sparse_blocks_to)), sparsity_layout_to
 
 
-@triton_op("blksprs::adapt_layout", mutates_args={})
+@triton_op("blksprs::adapt_layout_forward", mutates_args={})
 def adapt_layout_forward(x: Tensor,
                          sparsity_layout_from: Tensor, sparsity_reverse_lut_from: Tensor,
                          sparsity_block_size_from: int,
                          _: Tensor, sparsity_lut_to: Tensor,
                          sparsity_block_size_to: int,
                          n_sparse_blocks_to: int) -> Tensor:
-    output = torch.zeros(size=(n_sparse_blocks_to, sparsity_block_size_to, sparsity_block_size_to),
-                         dtype=x.dtype, device=x.device)
-
-    x_b, x_r, x_c = x.size()
-    x_b_s, x_r_s, x_c_s = stride(x)
-    s_l_x_b, s_l_x_r, s_l_x_c = sparsity_layout_from.size()
-    s_l_x_b_s, s_l_x_r_s, s_l_x_c_s = stride(sparsity_layout_from)
-    o_b, o_r, o_c = output.size()
-    o_b_s, o_r_s, o_c_s = stride(output)
-    s_lut_o_r, s_lut_o_c = sparsity_lut_to.size()
-    s_lut_o_r_s, s_lut_o_c_s = stride(sparsity_lut_to)
-
-    triton_grid = lambda meta: [o_b,
-                                triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
-
-    # TODO wrap
-    (adapt_layout_kernel[triton_grid]
-     (x,
-      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_reverse_lut_from,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_lut_to, s_lut_o_r, s_lut_o_r_s, s_lut_o_c_s,
-      sparsity_block_size_from,
-      sparsity_block_size_to))
-
-    return output
-
-
-def adapt_layout_backward(ctx, grad_output):
+    with torch.no_grad():
+        output = torch.zeros(size=(n_sparse_blocks_to, sparsity_block_size_to, sparsity_block_size_to),
+                             dtype=x.dtype, device=x.device)
+
+        x_b, x_r, x_c = x.size()
+        x_b_s, x_r_s, x_c_s = stride(x)
+        s_l_x_b, s_l_x_r, s_l_x_c = sparsity_layout_from.size()
+        s_l_x_b_s, s_l_x_r_s, s_l_x_c_s = stride(sparsity_layout_from)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
+        s_lut_o_r, s_lut_o_c = sparsity_lut_to.size()
+        s_lut_o_r_s, s_lut_o_c_s = stride(sparsity_lut_to)
+
+        triton_grid = lambda meta: [o_b,
+                                    triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
+
+        (wrap_triton(adapt_layout_kernel)[triton_grid]
+         (x,
+          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_reverse_lut_from,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_lut_to, s_lut_o_r, s_lut_o_r_s, s_lut_o_c_s,
+          sparsity_block_size_from,
+          sparsity_block_size_to))
+
+        return output
+
+
+def adapt_layout_wrapper_backward(ctx, grad_output):
     x, sparsity_layout_from, sparsity_layout_to = ctx.saved_tensors
     sparsity_block_size_from = ctx.sparsity_block_size_from
     sparsity_block_size_to = ctx.sparsity_block_size_to
@@ -501,4 +503,4 @@ def adapt_layout_setup_context(ctx, inputs, output):
     ctx.sparsity_block_size_to = sparsity_block_size_to
 
 
-adapt_layout_forward.register_autograd(adapt_layout_backward, setup_context=adapt_layout_setup_context)
+adapt_layout_forward.register_autograd(adapt_layout_wrapper_backward, setup_context=adapt_layout_setup_context)

blksprs/ops/distribution.py

@@ -51,44 +51,45 @@ def gather(src: BlksprsTensor, sparsity_layout_src: Tensor,
                                         sparsity_block_size))
 
 
-@triton_op("blksprs::gather", mutates_args={})
+@triton_op("blksprs::gather_forward", mutates_args={})
 def gather_forward(x: Tensor, sparsity_layout_x: Tensor, sparsity_reverse_lut_x: Tensor,
                    dim: int, i: Tensor, _: Tensor, sparsity_lut_i: Tensor,
                    sparsity_block_size: int) -> Tensor:
-    output = torch.zeros_like(i, dtype=x.dtype)
-
-    x_b, x_r, x_c = x.size()
-    x_b_s, x_r_s, x_c_s = stride(x)
-    s_l_x_b, s_l_x_r, s_l_x_c = sparsity_layout_x.size()
-    s_l_x_b_s, s_l_x_r_s, s_l_x_c_s = stride(sparsity_layout_x)
-    i_b, i_r, i_c = i.size()
-    i_b_s, i_r_s, i_c_s = stride(i)
-    s_lut_i_r, s_lut_i_c = sparsity_lut_i.size()
-    s_lut_i_r_s, s_lut_i_c_s = stride(sparsity_lut_i)
-    o_b, o_r, o_c = output.size()
-    o_b_s, o_r_s, o_c_s = stride(output)
-
-    triton_grid = lambda meta: [o_b,
-                                triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
-
-    (wrap_triton(gather_kernel)[triton_grid]
-     (x,
-      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_reverse_lut_x,
-      dim,
-      i,
-      i_b, i_b_s, i_r_s, i_c_s,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_lut_i, s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
-      sparsity_block_size))
-
-    return output
-
-
-def gather_backward(ctx, grad_output):
+    with torch.no_grad():
+        output = torch.zeros_like(i, dtype=x.dtype)
+
+        x_b, x_r, x_c = x.size()
+        x_b_s, x_r_s, x_c_s = stride(x)
+        s_l_x_b, s_l_x_r, s_l_x_c = sparsity_layout_x.size()
+        s_l_x_b_s, s_l_x_r_s, s_l_x_c_s = stride(sparsity_layout_x)
+        i_b, i_r, i_c = i.size()
+        i_b_s, i_r_s, i_c_s = stride(i)
+        s_lut_i_r, s_lut_i_c = sparsity_lut_i.size()
+        s_lut_i_r_s, s_lut_i_c_s = stride(sparsity_lut_i)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
+
+        triton_grid = lambda meta: [o_b,
+                                    triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
+
+        (wrap_triton(gather_kernel)[triton_grid]
+         (x,
+          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_reverse_lut_x,
+          dim,
+          i,
+          i_b, i_b_s, i_r_s, i_c_s,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_lut_i, s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
+          sparsity_block_size))
+
+        return output
+
+
+def gather_wrapper_backward(ctx, grad_output):
     sparsity_layout_x, i, sparsity_layout_i = ctx.saved_tensors
     dim = ctx.dim
     sparsity_block_size = ctx.sparsity_block_size
@@ -221,7 +222,7 @@ def gather_setup_context(ctx, inputs, output):
     ctx.sparsity_block_size = sparsity_block_size
 
 
-gather_forward.register_autograd(gather_backward, setup_context=gather_setup_context)
+gather_forward.register_autograd(gather_wrapper_backward, setup_context=gather_setup_context)
 
 
 def scatter(src: BlksprsTensor, sparsity_layout_src: Tensor,
@@ -288,52 +289,53 @@ def scatter_reduce(src: BlksprsTensor, sparsity_layout_src: Tensor,
                                                 reduce_op))
 
 
-@triton_op("blksprs::scatter_reduce", mutates_args={})
+@triton_op("blksprs::scatter_reduce_forward", mutates_args={})
 def scatter_reduce_forward(x: Tensor, _: Tensor, sparsity_lut_x: Tensor,
                            dim: int, i: Tensor,
                            sparsity_layout_o: Tensor, sparsity_reverse_lut_o: Tensor,
                            sparsity_block_size: int, n_sparse_blocks: int,
                            reduce_op: str) -> Tensor:
-    output = torch.zeros(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
-                         dtype=x.dtype, device=x.device)
-
-    x_b, x_r, x_c = x.size()
-    x_b_s, x_r_s, x_c_s = stride(x)
-    s_lut_x_r, s_lut_x_c = sparsity_lut_x.size()
-    s_lut_x_r_s, s_lut_x_c_s = stride(sparsity_lut_x)
-    i_b, i_r, i_c = i.size()
-    i_b_s, i_r_s, i_c_s = stride(i)
-    o_b, o_r, o_c = output.size()
-    o_b_s, o_r_s, o_c_s = stride(output)
-    s_l_o_b, s_l_o_r, s_l_o_c = sparsity_layout_o.size()
-    s_l_o_b_s, s_l_o_r_s, s_l_o_c_s = stride(sparsity_layout_o)
-
-    triton_grid = lambda meta: [x_b,
-                                triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
-
-    reduce_op_ind = 0
-    if reduce_op == "sum":
-        reduce_op_ind = 1
-
-    (wrap_triton(scatter_reduce_kernel)[triton_grid]
-     (x,
-      x_b, x_b_s, x_r_s, x_c_s,
-      sparsity_lut_x, s_lut_x_r, s_lut_x_r_s, s_lut_x_c_s,
-      dim,
-      i,
-      i_b, i_b_s, i_r_s, i_c_s,
-      output,
-      o_b, o_b_s,
-      s_l_o_b, s_l_o_b_s, s_l_o_r_s, s_l_o_c_s,
-      sparsity_reverse_lut_o,
-      reduce_op_ind,
-      sparsity_block_size))
-
-    return output
-
-
-def scatter_reduce_backward(ctx, grad_output):
+    with torch.no_grad():
+        output = torch.zeros(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
+                             dtype=x.dtype, device=x.device)
+
+        x_b, x_r, x_c = x.size()
+        x_b_s, x_r_s, x_c_s = stride(x)
+        s_lut_x_r, s_lut_x_c = sparsity_lut_x.size()
+        s_lut_x_r_s, s_lut_x_c_s = stride(sparsity_lut_x)
+        i_b, i_r, i_c = i.size()
+        i_b_s, i_r_s, i_c_s = stride(i)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
+        s_l_o_b, s_l_o_r, s_l_o_c = sparsity_layout_o.size()
+        s_l_o_b_s, s_l_o_r_s, s_l_o_c_s = stride(sparsity_layout_o)
+
+        triton_grid = lambda meta: [x_b,
+                                    triton.cdiv(x_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(x_c, meta["TRITON_BLOCK_SIZE"])]
+
+        reduce_op_ind = 0
+        if reduce_op == "sum":
+            reduce_op_ind = 1
+
+        (wrap_triton(scatter_reduce_kernel)[triton_grid]
+         (x,
+          x_b, x_b_s, x_r_s, x_c_s,
+          sparsity_lut_x, s_lut_x_r, s_lut_x_r_s, s_lut_x_c_s,
+          dim,
+          i,
+          i_b, i_b_s, i_r_s, i_c_s,
+          output,
+          o_b, o_b_s,
+          s_l_o_b, s_l_o_b_s, s_l_o_r_s, s_l_o_c_s,
+          sparsity_reverse_lut_o,
+          reduce_op_ind,
+          sparsity_block_size))
+
+        return output
+
+
+def scatter_reduce_wrapper_backward(ctx, grad_output):
     sparsity_layout_x, i, sparsity_layout_o = ctx.saved_tensors
     dim = ctx.dim
     sparsity_block_size = ctx.sparsity_block_size
@@ -477,4 +479,4 @@ def scatter_reduce_setup_context(ctx, inputs, output):
     ctx.reduce_op = reduce_op
 
 
-scatter_reduce_forward.register_autograd(scatter_reduce_backward, setup_context=scatter_reduce_setup_context)
+scatter_reduce_forward.register_autograd(scatter_reduce_wrapper_backward, setup_context=scatter_reduce_setup_context)