blksprs 2.1__py3-none-any.whl → 2.1.2__py3-none-any.whl

blksprs/__init__.py

@@ -1,6 +1,6 @@
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
-__version__ = "2.1"
+__version__ = "2.1.2"
 
 
 class ops:

blksprs/ops/conversion.py

@@ -13,7 +13,7 @@ from blksprs.utils.validation import validate_contiguous, validate_dimensions, v
 
 
 def to_blksprs(x: Tensor, sparsity_layout: Tensor, sparsity_block_size: int) -> BlksprsTensor:
-    """Wrapper for ``to_sparse``.
+    """Wrapper for :func:`to_sparse`.
 
     """
     return to_sparse(x, sparsity_layout, sparsity_block_size)
@@ -167,7 +167,7 @@ to_sparse_forward.register_autograd(to_sparse_wrapper_backward, setup_context=to
 
 def from_blksprs(x: BlksprsTensor, sparsity_layout: Tensor,
                  sparsity_block_size: int, fill_value: float = 0, lut: dict = None) -> Tensor:
-    """Wrapper for ``to_dense``.
+    """Wrapper for :func:`to_dense`.
 
     """
     return to_dense(x, sparsity_layout, sparsity_block_size, fill_value=fill_value, lut=lut)

blksprs/ops/softmax.py

@@ -9,15 +9,26 @@ from triton import language as tl
 
 from blksprs.ops.misc.row_wise import row_wise_sum, row_wise_max, row_wise_sub
 from blksprs.utils.blksprs_tensor import BlksprsTensor
-from blksprs.utils.debugging import dbg_tensor_full
-from blksprs.utils.tools import stride
+from blksprs.utils.tools import stride, ceil_pow2
 from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
     validate_sparsity, validate_sparsity_block_size, validate_dtype_float_32
 
 
+def softmax(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, flag_fused: bool = True,
+            lut: dict = None) -> BlksprsTensor:
+    """Wrapper for :func:`softmax_regular` and :func:`softmax_fused` based on the ``flag_fused`` parameter.
+
+    """
+    if flag_fused:
+        return softmax_fused(x, sparsity_layout, sparsity_block_size, lut)
+    else:
+        return softmax_regular(x, sparsity_layout, sparsity_block_size, lut)
+
+
 @torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float32)
-def softmax(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, lut: dict = None) -> BlksprsTensor:
+def softmax_regular(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int,
+                    lut: dict = None) -> BlksprsTensor:
     """Computes the softmax of a block-sparse tensor in compressed form.
 
     Note:
@@ -338,13 +349,15 @@ def softmax_fused(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size
     lut = softmax_fused_build_lut(lut, sparsity_layout)
 
     return BlksprsTensor(softmax_fused_forward(x, sparsity_layout,
-                                               lut["sparsity_reverse_lut"],
+                                               lut["sparsity_reverse_lut_sorted"],
+                                               lut["max_blocks_line"],
                                                sparsity_block_size))
 
 
 @triton_op("blksprs::softmax_fused_forward", mutates_args={})
 def softmax_fused_forward(x: Tensor, sparsity_layout: Tensor,
-                          sparsity_reverse_lut: Tensor,
+                          sparsity_reverse_lut_sorted: Tensor,
+                          max_blocks_line: int,
                           sparsity_block_size: int) -> Tensor:
     output = torch.zeros_like(x)
 
@@ -361,23 +374,29 @@ def softmax_fused_forward(x: Tensor, sparsity_layout: Tensor,
      (x,
       x_b, x_b_s, x_r_s, x_c_s,
       output,
-      s_l_b, s_l_b_s, s_l_r_s, s_l_c, s_l_c_s,
-      sparsity_reverse_lut,
+      s_l_b, s_l_b_s, s_l_r_s, s_l_c_s,
+      sparsity_reverse_lut_sorted,
+      max_blocks_line,
       sparsity_block_size))
 
     return output
 
 
 def softmax_fused_backward_wrapper(ctx, grad_output):
-    o, sparsity_layout, sparsity_reverse_lut = ctx.saved_tensors
+    o, sparsity_layout, sparsity_reverse_lut_sorted = ctx.saved_tensors
+    max_blocks_line = ctx.max_blocks_line
     sparsity_block_size = ctx.sparsity_block_size
 
-    return softmax_fused_backward(grad_output, o, sparsity_reverse_lut, sparsity_layout,
-                                  sparsity_block_size), None, None, None, None, None
+    return softmax_fused_backward(grad_output, o, sparsity_reverse_lut_sorted, sparsity_layout,
+                                  max_blocks_line, sparsity_block_size), None, None, None, None
 
 
 @triton_op("blksprs::softmax_fused_backward", mutates_args={})
-def softmax_fused_backward(grad_output: Tensor, o: Tensor, sparsity_reverse_lut: Tensor, sparsity_layout: Tensor,
+def softmax_fused_backward(grad_output: Tensor,
+                           o: Tensor,
+                           sparsity_reverse_lut_sorted: Tensor,
+                           sparsity_layout: Tensor,
+                           max_blocks_line: int,
                            sparsity_block_size: int) -> Tensor:
     with torch.no_grad():
         grad_x = torch.zeros_like(o)
@@ -398,9 +417,10 @@ def softmax_fused_backward(grad_output: Tensor, o: Tensor, sparsity_reverse_lut:
           g_b, g_b_s, g_r_s, g_c_s,
           o,
           o_b, o_b_s, o_r_s, o_c_s,
-          s_l_b, s_l_b_s, s_l_r_s, s_l_c, s_l_c_s,
-          sparsity_reverse_lut,
+          s_l_b, s_l_b_s, s_l_r_s, s_l_c_s,
+          sparsity_reverse_lut_sorted,
           grad_x,
+          max_blocks_line,
           sparsity_block_size))
 
         return grad_x
@@ -417,8 +437,9 @@ def softmax_fused_backward(grad_output: Tensor, o: Tensor, sparsity_reverse_lut:
 def softmax_fused_kernel(x,
                          x_b, x_b_s, x_r_s, x_c_s,
                          o,
-                         s_l_b, s_l_b_s, s_l_r_s, s_l_c: tl.constexpr, s_l_c_s,
-                         r_lut,
+                         s_l_b, s_l_b_s, s_l_r_s, s_l_c_s,
+                         r_lut_s,
+                         mbs: tl.constexpr,
                          sparsity_block_size: tl.constexpr,
                          TRITON_BLOCK_SIZE: tl.constexpr) -> None:
     # Get triton block indices
@@ -429,21 +450,21 @@ def softmax_fused_kernel(x,
     # Load reverse sparsity indices of row
     blk_rev_idx = (pid_bat * s_l_b_s +
                    pid_row * s_l_r_s +
-                   (tl.arange(0, s_l_c) * s_l_c_s))
+                   (tl.arange(0, mbs) * s_l_c_s))
     blk_rev_msk = (blk_rev_idx >= 0 and blk_rev_idx < s_l_b * s_l_b_s)
-    blk_rev = tl.load(r_lut + blk_rev_idx, mask=blk_rev_msk).to(tl.int32)
+    blk_rev = tl.load(r_lut_s + blk_rev_idx, mask=blk_rev_msk).to(tl.int32)
 
     if (not (tl.min(blk_rev) == -1 and
              tl.max(blk_rev) == -1)):
         # Extend sparsity indices to cover sparsity blocks
         blk_rev_ext = tl.expand_dims(blk_rev, -1)
-        blk_rev_ext = tl.broadcast_to(blk_rev_ext, (s_l_c, sparsity_block_size))
-        blk_rev_ext = tl.reshape(blk_rev_ext, (s_l_c * sparsity_block_size))
+        blk_rev_ext = tl.broadcast_to(blk_rev_ext, (mbs, sparsity_block_size))
+        blk_rev_ext = tl.reshape(blk_rev_ext, (mbs * sparsity_block_size))
 
         # Load line of x
         blk_x_idx = (blk_rev_ext * x_b_s +
                      pid_lin * x_r_s +
-                     (tl.arange(0, s_l_c * sparsity_block_size) % sparsity_block_size) * x_c_s)
+                     (tl.arange(0, mbs * sparsity_block_size) % sparsity_block_size) * x_c_s)
         blk_x_mask = ((blk_x_idx >= 0 and blk_x_idx < x_b * x_b_s)
                       and blk_rev_ext != -1)
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_mask, other=float("-inf"))
@@ -467,9 +488,10 @@ def softmax_fused_kernel_grad(g,
                               g_b, g_b_s, g_r_s, g_c_s,
                               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: tl.constexpr, s_l_c_s,
-                              r_lut,
+                              s_l_b, s_l_b_s, s_l_r_s, s_l_c_s,
+                              r_lut_s,
                               o,
+                              mbs: tl.constexpr,
                               sparsity_block_size: tl.constexpr,
                               TRITON_BLOCK_SIZE: tl.constexpr) -> None:
     # Get triton block indices
@@ -480,21 +502,21 @@ def softmax_fused_kernel_grad(g,
     # Load reverse sparsity indices of row
     blk_rev_idx = (pid_bat * s_l_b_s +
                    pid_row * s_l_r_s +
-                   (tl.arange(0, s_l_c) * s_l_c_s))
+                   (tl.arange(0, mbs) * s_l_c_s))
     blk_rev_msk = (blk_rev_idx >= 0 and blk_rev_idx < s_l_b * s_l_b_s)
-    blk_rev = tl.load(r_lut + blk_rev_idx, mask=blk_rev_msk).to(tl.int32)
+    blk_rev = tl.load(r_lut_s + blk_rev_idx, mask=blk_rev_msk).to(tl.int32)
 
     if (not (tl.min(blk_rev) == -1 and
              tl.max(blk_rev) == -1)):
         # Extend sparsity indices to cover sparsity blocks
         blk_rev_ext = tl.expand_dims(blk_rev, -1)
-        blk_rev_ext = tl.broadcast_to(blk_rev_ext, (s_l_c, sparsity_block_size))
-        blk_rev_ext = tl.reshape(blk_rev_ext, (s_l_c * sparsity_block_size))
+        blk_rev_ext = tl.broadcast_to(blk_rev_ext, (mbs, sparsity_block_size))
+        blk_rev_ext = tl.reshape(blk_rev_ext, (mbs * sparsity_block_size))
 
         # Load line of g
         blk_g_idx = (blk_rev_ext * g_b_s +
                      pid_lin * g_r_s +
-                     (tl.arange(0, s_l_c * sparsity_block_size) % sparsity_block_size) * g_c_s)
+                     (tl.arange(0, mbs * sparsity_block_size) % sparsity_block_size) * g_c_s)
         blk_g_mask = ((blk_g_idx >= 0 and blk_g_idx < g_b * g_b_s)
                       and blk_rev_ext != -1)
         blk_g = tl.load(g + blk_g_idx, mask=blk_g_mask)
@@ -502,7 +524,7 @@ def softmax_fused_kernel_grad(g,
         # Load line of x
         blk_x_idx = (blk_rev_ext * x_b_s +
                      pid_lin * x_r_s +
-                     (tl.arange(0, s_l_c * sparsity_block_size) % sparsity_block_size) * x_c_s)
+                     (tl.arange(0, mbs * sparsity_block_size) % sparsity_block_size) * x_c_s)
         blk_x_mask = ((blk_x_idx >= 0 and blk_x_idx < x_b * x_b_s)
                       and blk_rev_ext != -1)
         blk_x = tl.load(x + blk_x_idx, mask=blk_x_mask)
@@ -510,6 +532,7 @@ def softmax_fused_kernel_grad(g,
         # Compute gradients
         blk_grad = blk_x * (blk_g - tl.sum(blk_x * blk_g))
 
+        # Store output
         tl.store(o + blk_x_idx, blk_grad, mask=blk_x_mask)
 
 
@@ -517,25 +540,36 @@ def softmax_fused_build_lut(lut: dict, sparsity_layout: Tensor):
     if lut is None:
         lut = dict()
 
-    if "sparsity_reverse_lut" not in lut:
+    if "sparsity_reverse_lut_sorted" 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())
-                                .reshape(-1).contiguous())
-        lut["sparsity_reverse_lut"] = sparsity_reverse_lut
-
-    validate_contiguous(sparsity_layout, lut["sparsity_reverse_lut"])
+        sparsity_reverse_lut_sorted = (((torch.cumsum(sparsity_layout_flat, dim=-1) - 1) *
+                                        (sparsity_layout_flat == 1) -
+                                        (1 * (sparsity_layout_flat == 0)))
+                                       .reshape(sparsity_layout.size())
+                                       .sort(descending=True, dim=-1)[0]
+                                       .reshape(-1).contiguous())
+        lut["sparsity_reverse_lut_sorted"] = sparsity_reverse_lut_sorted
+
+    if "max_blocks_line" not in lut:
+        sparsity_reverse_lut_sorted = lut["sparsity_reverse_lut_sorted"]
+        max_blocks_line = ((torch.reshape(sparsity_reverse_lut_sorted, (-1, sparsity_layout.size(-1)))
+                            != -1)
+                           .sum(dim=-1)
+                           .max()
+                           .item())
+        lut["max_blocks_line"] = min(ceil_pow2(max(max_blocks_line, 2)), sparsity_layout.size(-1))
+
+    validate_contiguous(sparsity_layout, lut["sparsity_reverse_lut_sorted"])
 
     return lut
 
 
 # noinspection PyUnusedLocal
 def softmax_fused_setup_context(ctx, inputs, output):
-    (_, sparsity_layout, sparsity_reverse_lut, sparsity_block_size) = inputs
+    (_, sparsity_layout, sparsity_reverse_lut_sorted, max_blocks_line, sparsity_block_size) = inputs
 
-    ctx.save_for_backward(output, sparsity_layout, sparsity_reverse_lut)
+    ctx.save_for_backward(output, sparsity_layout, sparsity_reverse_lut_sorted)
+    ctx.max_blocks_line = max_blocks_line
     ctx.sparsity_block_size = sparsity_block_size
 
 

blksprs/utils/tools.py

@@ -26,3 +26,8 @@ def stride(x: Tensor):
         return x.size(1) * x.size(2), x.size(2), 1
     else:
         raise NotImplementedError
+
+def ceil_pow2(x: int) -> int:
+    if x <= 0:
+        raise ValueError("Input must be a positive integer.")
+    return 1 << (x - 1).bit_length()

{blksprs-2.1.dist-info → blksprs-2.1.2.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.1
-Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
+Version: 2.1.2
+Summary: A lightweight library for operations on block-sparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs
 Project-URL: Bugtracker, https://github.com/FelixSchoen/blksprs/issues
@@ -197,6 +197,7 @@ def test_readme():
     # Other available functions
     bs.ops.transpose(o_sparse, sparsity_layout_o, sparsity_block_size)
     bs.ops.softmax(o_sparse, sparsity_layout_o, sparsity_block_size)
+    bs.ops.softmax_fused(o_sparse, sparsity_layout_o, sparsity_block_size) # Significantly faster version that requires that rows of matrix fit into memory
     bs.ops.misc.row_wise_sum(o_sparse, sparsity_layout_o, sparsity_block_size)
     bs.ops.misc.row_wise_max(o_sparse, sparsity_layout_o, sparsity_block_size)
 

{blksprs-2.1.dist-info → blksprs-2.1.2.dist-info}/RECORD

@@ -1,13 +1,13 @@
-blksprs/__init__.py,sha256=o_Rj7fz_70vbMGLePihczVIVcM8E28vY3ah-d1q4ZO0,1613
+blksprs/__init__.py,sha256=NRxydw4i9jg7WeDuojfEePdtdbughV9AZsEcT9yywK4,1615
 blksprs/layouting/distribution_layout.py,sha256=ur1ty_2U-Hfj78hMWsLZvu7ZuGhzW3qGLKMc72DfTZM,5861
 blksprs/layouting/sparsity_layout.py,sha256=eXHmu2h7K5Q-YUpfOxocJoeP_5ZoQFZf_eHLxRZQbYU,11207
-blksprs/ops/conversion.py,sha256=kf5HKofZ4nVeHCIqQoYKiIlgsAhq33Tnmnr1c17Fkqs,21906
+blksprs/ops/conversion.py,sha256=RgVSyiULLwv8KWQqSyXpKwTr4Qp-lpDK9i-zKlN841I,21914
 blksprs/ops/distribution.py,sha256=0tPldv0ARzmCV1CU2jvfqpHBgOuHPrDFiCtqsLs7CZc,20789
 blksprs/ops/flow.py,sha256=qdWBCLDSkKaa8CAfkO1NgH-J5N7yMsILyR7qEpyrIUU,8246
 blksprs/ops/matmul.py,sha256=5tVBKU_lglUjaLDi6J_dscdqlmzRz38OGxqAxZxZXDs,11879
 blksprs/ops/partitioning.py,sha256=cfQmY9BZqGTvvJorIhtb-EyuGRJGPraWR-wTKdb47aI,9954
 blksprs/ops/repeat.py,sha256=TLYNxwPuT9y5K9xyM41WK5gnggAJF3lI61Q2K7zWjns,9035
-blksprs/ops/softmax.py,sha256=H0OxST_XX1QLa7HDTDHznzibVHAxnp5sVbMU32HLxf0,21967
+blksprs/ops/softmax.py,sha256=1UIovPrdE_zgAIPqjmOTFn8CMbd_2Z8tPP-vMBxU07I,23526
 blksprs/ops/transpose.py,sha256=U-VAyLRT6_NDv9qYSFzBqfVlDeIpTqAMEXkqto0VF6w,4072
 blksprs/ops/misc/broadcast_ops.py,sha256=-PrHiSJikZh8nXUmXxSCtFEP27TTxFr4wcrNxBjnimk,5987
 blksprs/ops/misc/row_wise.py,sha256=n5FJjAuOd8BHBJQx4bsQwr-HmXkR9PYVAqfk77wjOFU,19653
@@ -15,9 +15,9 @@ blksprs/utils/autotuning.py,sha256=a-kmWRjJ3eED2XbjkQeOJSyW8bdIs27HgKMPvAKqWeU,2
 blksprs/utils/benchmarking.py,sha256=dLabDscTFn5NkmOI1g7DnKeTneUYW3RIVv9MDF-8BKc,1271
 blksprs/utils/blksprs_tensor.py,sha256=pfoz59aJixj_fIoFx76ySiygwRQUemmgjMKepZ2c4j0,244
 blksprs/utils/processing.py,sha256=RNkEDc0g-sNHRuMPkRzNWU13d3_lIkXMJdoqES4yQTM,3738
-blksprs/utils/tools.py,sha256=CPf7viQ2OTcZFrB1aSL8_us4VE9M6YEfDz2dE30jr9I,715
+blksprs/utils/tools.py,sha256=TKygEKge4wJtJnXXDg8BTL8vzBpqIJsQ_A3_5FmLpcE,859
 blksprs/utils/validation.py,sha256=G8eQlvJVMKfEX3k2AwBD0A6Ck-gFoRLpLNY6HXsB3fA,4348
-blksprs-2.1.dist-info/METADATA,sha256=uPVm8Y7fX5iModz6j3hNAftdtauCsJ-iYrMa-Pv3xnU,9506
-blksprs-2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-blksprs-2.1.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-2.1.dist-info/RECORD,,
+blksprs-2.1.2.dist-info/METADATA,sha256=U20ZL7XLhrgiMd_0QGFik0Ci43SDoCT8q876-1yCeNo,9665
+blksprs-2.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+blksprs-2.1.2.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-2.1.2.dist-info/RECORD,,