blksprs 2.0rc8__py3-none-any.whl → 2.1__py3-none-any.whl

blksprs/__init__.py

@@ -1,12 +1,13 @@
-from blksprs.utils.tools import version
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
+__version__ = "2.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/ops/conversion.py

@@ -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/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/ops/softmax.py

@@ -1,3 +1,5 @@
+import pdb
+
 import torch
 import triton
 from torch import Tensor
@@ -7,6 +9,7 @@ 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.autotuning import get_autotune_configs, prune_autotune_configs
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \
@@ -100,6 +103,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 +121,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 +305,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/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.dist-info → blksprs-2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.0rc8
+Version: 2.1
 Summary: A lightweight library for operations on blocksparse matrices in PyTorch.
 Author-email: Felix Schön <schoen@kr.tuwien.ac.at>
 Project-URL: Homepage, https://github.com/FelixSchoen/blksprs

{blksprs-2.0rc8.dist-info → blksprs-2.1.dist-info}/RECORD

@@ -1,13 +1,13 @@
-blksprs/__init__.py,sha256=283rF0fbrUqsH_KXUvCgbCMqO0GOgenMkxwDVh1QdpU,1617
+blksprs/__init__.py,sha256=o_Rj7fz_70vbMGLePihczVIVcM8E28vY3ah-d1q4ZO0,1613
 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=_g32aEEZdeuHHPj1pBfTNMxknRwJ9O1zk3Wv76pBIrg,21898
+blksprs/ops/conversion.py,sha256=kf5HKofZ4nVeHCIqQoYKiIlgsAhq33Tnmnr1c17Fkqs,21906
 blksprs/ops/distribution.py,sha256=0tPldv0ARzmCV1CU2jvfqpHBgOuHPrDFiCtqsLs7CZc,20789
-blksprs/ops/flow.py,sha256=PDZAD8u4y9qW1IXERki6ItKbEKnm_ChG8SKWM3_P9Oc,8245
+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=BwrRQdtRdkiSvl2mf5bpsTmyIxWiJOpa1HFg0st5yGU,12778
+blksprs/ops/softmax.py,sha256=H0OxST_XX1QLa7HDTDHznzibVHAxnp5sVbMU32HLxf0,21967
 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=BozpH3oEXe3K9ZRJsIzlasDk-sZyJqmwSf1gl7xbbdo,865
+blksprs/utils/tools.py,sha256=CPf7viQ2OTcZFrB1aSL8_us4VE9M6YEfDz2dE30jr9I,715
 blksprs/utils/validation.py,sha256=G8eQlvJVMKfEX3k2AwBD0A6Ck-gFoRLpLNY6HXsB3fA,4348
-blksprs-2.0rc8.dist-info/METADATA,sha256=h70L26BthR6laP7sMQLF9L3dHIRQNCF_oKwZ5g4dZSg,9509
-blksprs-2.0rc8.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-blksprs-2.0rc8.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-2.0rc8.dist-info/RECORD,,
+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.0rc8.dist-info → blksprs-2.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.8.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any