blksprs 2.0rc8__tar.gz → 2.1.1__tar.gz

{blksprs-2.0rc8 → blksprs-2.1.1}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.0rc8
-Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
+Version: 2.1.1
+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.0rc8 → blksprs-2.1.1}/README.md

@@ -178,6 +178,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.0rc8 → blksprs-2.1.1}/blksprs/__init__.py

@@ -1,12 +1,13 @@
-from blksprs.utils.tools import version
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
+__version__ = "2.1.1"
+
 
 class ops:
     from blksprs.ops.conversion import to_dense, to_sparse, from_blksprs, to_blksprs, adapt_layout
     from blksprs.ops.distribution import gather, scatter, scatter_reduce
     from blksprs.ops.matmul import matmul
-    from blksprs.ops.softmax import softmax
+    from blksprs.ops.softmax import softmax, softmax_fused
     from blksprs.ops.transpose import transpose
     from blksprs.ops.repeat import repeat, repeat_interleave
     from blksprs.ops.partitioning import split, merge

{blksprs-2.0rc8 → blksprs-2.1.1}/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)
@@ -204,8 +204,8 @@ def to_dense(x: BlksprsTensor, sparsity_layout: Tensor,
     if sparsity_layout.size(1) == 1 and sparsity_layout.size(2) == 1 and torch.all(sparsity_layout):
         return x
 
-    return to_dense_forward(x, sparsity_layout,
-                            lut["sparsity_reverse_lut"], sparsity_block_size, fill_value)
+    return Tensor(to_dense_forward(x, sparsity_layout,
+                            lut["sparsity_reverse_lut"], sparsity_block_size, fill_value))
 
 
 @triton_op("blksprs::to_dense_forward", mutates_args={})

{blksprs-2.0rc8 → blksprs-2.1.1}/blksprs/ops/flow.py

@@ -78,7 +78,7 @@ def flow_pull_kernel(x,
     spa_col_msk = (spa_col_idx >= 0 and spa_col_idx < s_lut_r * s_lut_r_s)
     spa_col = tl.load(s_lut + spa_col_idx, mask=spa_col_msk)
 
-    # Get reverse sparsity index
+    # Load reverse sparsity index
     rev_idx_spa_idx = (spa_bat * s_l_o_b_s +
                        spa_row * s_l_o_r_s +
                        spa_col * s_l_o_c_s)

{blksprs-2.0rc8 → blksprs-2.1.1}/blksprs/ops/softmax.py

@@ -1,3 +1,5 @@
+import pdb
+
 import torch
 import triton
 from torch import Tensor
@@ -13,8 +15,20 @@ from blksprs.utils.validation import validate_contiguous, validate_dimensions, v
     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:
@@ -100,6 +114,8 @@ def softmax_backward_wrapper(ctx, grad_output):
 def softmax_backward(grad_output: Tensor, o: Tensor, sparsity_lut: Tensor, sparsity_layout: Tensor,
                      sparsity_block_size: int) -> Tensor:
     with torch.no_grad():
+        grad_x = torch.zeros_like(o, dtype=torch.float)
+
         s, sparsity_layout_s = row_wise_sum(grad_output * o, sparsity_layout, sparsity_block_size, flag_slice_only=True)
 
         sparsity_layout_s_flat = sparsity_layout_s.reshape(-1)
@@ -116,8 +132,6 @@ def softmax_backward(grad_output: Tensor, o: Tensor, sparsity_lut: Tensor, spars
         s_l_s_b, s_l_s_r, s_l_s_c = sparsity_layout_s.size()
         s_l_s_b_s, s_l_s_r_s, s_l_s_c_s = stride(sparsity_layout_s)
 
-        grad_x = torch.zeros_like(o, dtype=torch.float)
-
         triton_grid = lambda meta: [o_b,
                                     triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
                                     triton.cdiv(o_c, meta["TRITON_BLOCK_SIZE"])]
@@ -302,3 +316,238 @@ def softmax_setup_context(ctx, inputs, output):
 
 
 softmax_forward.register_autograd(softmax_backward_wrapper, setup_context=softmax_setup_context)
+
+
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float32)
+def softmax_fused(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int,
+                  lut: dict = None) -> BlksprsTensor:
+    """Computes the softmax fused for each row of a block-sparse tensor in compressed form.
+
+    Note:
+        This softmax implementation is a fused version that loads the entire row of a block-sparse tensor into memory.
+        See :func:`softmax` for a true block-wise softmax implementation.
+
+    Args:
+        x (BlksprsTensor): A block-sparse tensor in compressed form.
+        sparsity_layout (Tensor): The sparsity layout of the block-sparse tensor.
+        sparsity_block_size (int): The size of the sparsity blocks.
+        lut (dict, optional): A dictionary containing the look-up tables for the operation (default ``None``).
+
+    Returns:
+        BlksprsTensor: The result of the softmax operation as a block-sparse tensor in compressed form.
+
+    """
+    x = x.contiguous()
+
+    validate_dimensions(x)
+    validate_contiguous(x)
+    validate_dtype_float_32(x)
+    validate_device(x)
+    validate_sparsity(sparsity_block_size, (x, sparsity_layout))
+    validate_sparsity_block_size(sparsity_block_size, x)
+
+    lut = softmax_fused_build_lut(lut, sparsity_layout)
+
+    return BlksprsTensor(softmax_fused_forward(x, sparsity_layout,
+                                               lut["sparsity_reverse_lut"],
+                                               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_block_size: int) -> Tensor:
+    output = torch.zeros_like(x)
+
+    x_b, x_r, x_c = x.size()
+    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)
+
+    triton_grid = lambda meta: [s_l_b,
+                                s_l_r,
+                                sparsity_block_size]
+
+    (wrap_triton(softmax_fused_kernel)[triton_grid]
+     (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,
+      sparsity_block_size))
+
+    return output
+
+
+def softmax_fused_backward_wrapper(ctx, grad_output):
+    o, sparsity_layout, sparsity_reverse_lut = ctx.saved_tensors
+    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
+
+
+@triton_op("blksprs::softmax_fused_backward", mutates_args={})
+def softmax_fused_backward(grad_output: Tensor, o: Tensor, sparsity_reverse_lut: Tensor, sparsity_layout: Tensor,
+                           sparsity_block_size: int) -> Tensor:
+    with torch.no_grad():
+        grad_x = torch.zeros_like(o)
+
+        g_b, g_r, g_c = grad_output.size()
+        g_b_s, g_r_s, g_c_s = stride(grad_output)
+        o_b, o_r, o_c = o.size()
+        o_b_s, o_r_s, o_c_s = stride(o)
+        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)
+
+        triton_grid = lambda meta: [s_l_b,
+                                    s_l_r,
+                                    sparsity_block_size]
+
+        (wrap_triton(softmax_fused_kernel_grad)[triton_grid]
+         (grad_output,
+          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,
+          grad_x,
+          sparsity_block_size))
+
+        return grad_x
+
+
+# noinspection PyUnusedLocal
+@triton.autotune(
+    configs=get_autotune_configs(),
+    key=["sparsity_block_size"],
+    prune_configs_by={"early_config_prune": prune_autotune_configs},
+    reset_to_zero=["o"]
+)
+@triton.jit
+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,
+                         sparsity_block_size: tl.constexpr,
+                         TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+    # Get triton block indices
+    pid_bat = tl.program_id(axis=0)
+    pid_row = tl.program_id(axis=1)
+    pid_lin = tl.program_id(axis=2)
+
+    # 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))
+    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)
+
+    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))
+
+        # 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)
+        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"))
+
+        # Compute softmax
+        blk_x_softmax = tl.softmax(blk_x)
+
+        # Store output
+        tl.store(o + blk_x_idx, blk_x_softmax, mask=blk_x_mask)
+
+
+# noinspection PyUnusedLocal
+@triton.autotune(
+    configs=get_autotune_configs(),
+    key=["sparsity_block_size"],
+    prune_configs_by={"early_config_prune": prune_autotune_configs},
+    reset_to_zero=["o"]
+)
+@triton.jit
+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,
+                              o,
+                              sparsity_block_size: tl.constexpr,
+                              TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+    # Get triton block indices
+    pid_bat = tl.program_id(axis=0)
+    pid_row = tl.program_id(axis=1)
+    pid_lin = tl.program_id(axis=2)
+
+    # 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))
+    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)
+
+    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))
+
+        # 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)
+        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)
+
+        # 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)
+        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)
+
+        # Compute gradients
+        blk_grad = blk_x * (blk_g - tl.sum(blk_x * blk_g))
+
+        tl.store(o + blk_x_idx, blk_grad, mask=blk_x_mask)
+
+
+def softmax_fused_build_lut(lut: dict, sparsity_layout: Tensor):
+    if lut is None:
+        lut = dict()
+
+    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())
+                                .reshape(-1).contiguous())
+        lut["sparsity_reverse_lut"] = sparsity_reverse_lut
+
+    validate_contiguous(sparsity_layout, lut["sparsity_reverse_lut"])
+
+    return lut
+
+
+# noinspection PyUnusedLocal
+def softmax_fused_setup_context(ctx, inputs, output):
+    (_, sparsity_layout, sparsity_reverse_lut, sparsity_block_size) = inputs
+
+    ctx.save_for_backward(output, sparsity_layout, sparsity_reverse_lut)
+    ctx.sparsity_block_size = sparsity_block_size
+
+
+softmax_fused_forward.register_autograd(softmax_fused_backward_wrapper, setup_context=softmax_fused_setup_context)

{blksprs-2.0rc8 → blksprs-2.1.1}/blksprs/utils/tools.py

@@ -1,6 +1,3 @@
-import tomllib
-from pathlib import Path
-
 import torch
 from torch import Tensor, Size
 
@@ -8,19 +5,14 @@ from torch import Tensor, Size
 torch._dynamo.config.capture_scalar_outputs = True
 
 
-def version():
-    with open(Path(__file__).parent.parent.parent.joinpath("pyproject.toml"), "rb") as f:
-        return tomllib.load(f)["project"]["version"]
-
-
-def do_shape_blocksparse(x: Tensor):
+def do_shape_blocksparse(x: Tensor) -> tuple[Tensor, Size]:
     if x.dim() == 3:
         return x.contiguous(), x.size()
 
     return x.reshape(-1, x.size(-2), x.size(-1)).contiguous(), x.size()
 
 
-def undo_shape_blocksparse(x: Tensor, shape: Size):
+def undo_shape_blocksparse(x: Tensor, shape: Size | tuple[int, ...]) -> Tensor:
     if x.shape[:-2] == shape[:-2]:
         return x
 

{blksprs-2.0rc8 → blksprs-2.1.1}/blksprs.egg-info/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.0rc8
-Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
+Version: 2.1.1
+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.0rc8 → blksprs-2.1.1}/pyproject.toml

@@ -1,8 +1,8 @@
 [project]
 name = "blksprs"
-version = "2.0-rc.8"
+version = "2.1.1"
 authors = [{ name = "Felix Schön", email = "schoen@kr.tuwien.ac.at" }]
-description = "A lightweight library for operations on blocksparse matrices in PyTorch."
+description = "A lightweight library for operations on block-sparse matrices in PyTorch."
 readme = "README.md"
 requires-python = ">=3.11"
 license = { file = "LICENSE.md" }