blksprs 1.8.2__py3-none-any.whl → 1.8.3__py3-none-any.whl

blksprs/__init__.py

@@ -1,27 +1,40 @@
-from blksprs.ops.conversion import to_dense, to_sparse, from_blksprs, to_blksprs
-from blksprs.ops.distribution import gather, scatter, scatter_reduce
-from blksprs.ops.matmul import matmul
-from blksprs.ops.softmax import softmax
-from blksprs.ops.transpose import transpose
-from blksprs.ops.repeat import repeat, repeat_interleave
-from blksprs.misc.partitioning import split, merge
+from blksprs.utils.blksprs_tensor import BlksprsTensor
 
+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.transpose import transpose
+    from blksprs.ops.repeat import repeat, repeat_interleave
+    from blksprs.ops.partitioning import split, merge
 
-class layout:
+    class misc:
+        from blksprs.ops.misc.row_wise import row_wise_sum, row_wise_max, row_wise_add, row_wise_sub
+        from blksprs.ops.misc.broadcast_ops import broadcast_add, broadcast_sub
+        from blksprs.ops.misc.exp import exp
+
+    class experimental:
+        from blksprs.ops.experimental.distribution_mdi import gather_mdi, scatter_reduce_mdi
+
+
+class layouting:
     from blksprs.layouting.distribution_layout import build_distribution_layout
     from blksprs.layouting.sparsity_layout import build_sparsity_layout, build_sparsity_layout_adaption, \
         build_sparsity_layout_matmul, build_sparsity_layout_matmul_fast
 
+    class experimental:
+        from blksprs.ops.experimental.distribution_mdi import build_distribution_layout_mdi
 
-class misc:
-    from blksprs.misc.broadcast_ops import broadcast_add, broadcast_sub
-    from blksprs.misc.exp import exp
-    from blksprs.misc.row_wise import row_wise_sum, row_wise_max, row_wise_add, row_wise_sub
-
-
-class util:
-    from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse, disable_validation
 
+class utils:
+    from blksprs.utils.processing import apply_torch_linear
+    from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse
+    from blksprs.utils.validation import disable_validation
 
-class experimental:
-    from blksprs.experimental.distribution_mdi import gather_mdi, scatter_reduce_mdi
+    class validation:
+        from blksprs.utils.validation import disable_validation
+        from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_dtype_float, \
+            validate_dtype_int, validate_device, validate_sparsity, validate_sparsity_dense, \
+            validate_sparsity_block_size, \
+            validate_triton_block_size

blksprs/ops/softmax.py

@@ -3,8 +3,8 @@ import triton
 from torch import Tensor
 from triton import language as tl
 
-from blksprs.misc.exp import exp
-from blksprs.misc.row_wise import row_wise_sum, row_wise_max, row_wise_sub
+from blksprs.ops.misc.exp import exp
+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.tools import get_triton_block_size, stride
 from blksprs.utils.validation import validate_contiguous, validate_dimensions, validate_device, \

blksprs/utils/processing.py

@@ -0,0 +1,41 @@
+import torch
+from torch import Tensor, nn
+from triton.language import dtype
+
+from blksprs.layouting.sparsity_layout import build_sparsity_layout_matmul_fast
+from blksprs.ops.conversion import to_sparse
+from blksprs.ops.matmul import matmul
+from blksprs.ops.repeat import repeat
+from blksprs.utils.blksprs_tensor import BlksprsTensor
+
+
+def apply_torch_linear(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int,
+                       linear: nn.Linear) -> (BlksprsTensor, Tensor):
+    # Extract weight and bias
+    w = linear.weight
+    b = linear.bias
+
+    # Convert w to block-sparse representation
+    sparsity_layout_w_t = torch.ones(size=(sparsity_layout.size(0), w.size(1) // sparsity_block_size,
+                                           w.size(0) // sparsity_block_size), dtype=torch.bool, device=x.device)
+    w_t_bs = to_sparse(w.transpose(-1, -2).unsqueeze(0).repeat(sparsity_layout.size(0), 1, 1),
+                       sparsity_layout_w_t, sparsity_block_size)
+
+    # Apply weights
+    sparsity_layout_xw = build_sparsity_layout_matmul_fast(sparsity_layout, sparsity_layout_w_t)
+    xw = matmul(x, sparsity_layout, w_t_bs, sparsity_layout_w_t, sparsity_layout_xw, sparsity_block_size)
+    interim = xw
+
+    # Apply bias
+    if b is not None:
+        b_slice = b.unsqueeze(0).unsqueeze(0).repeat(1, sparsity_block_size, 1)
+        sparsity_layout_b_slice = torch.ones(size=(1, b_slice.size(1) // sparsity_block_size,
+                                                   b_slice.size(2) // sparsity_block_size), dtype=torch.bool,
+                                             device=x.device)
+        b_slice_bs = to_sparse(b_slice, sparsity_layout_b_slice, sparsity_block_size)
+        b_bs, sparsity_layout_b = repeat(b_slice_bs, sparsity_layout_b_slice,
+                                         (sparsity_layout.size(0), sparsity_layout_xw.size(1), 1), sparsity_block_size,
+                                         sparsity_layout_output=sparsity_layout_xw)
+        interim = interim + b_bs
+
+    return interim, sparsity_layout_xw

blksprs/utils/tools.py

@@ -1,7 +1,5 @@
 from torch import Tensor, Size
 
-from blksprs.utils.validation import _set_skip_validation
-
 
 def do_shape_blocksparse(x: Tensor):
     if x.dim() == 3:
@@ -21,8 +19,5 @@ def get_triton_block_size(sparsity_block_size: int, limit: int = 128):
     return min(sparsity_block_size, limit)
 
 
-def disable_validation():
-    _set_skip_validation(True)
-
 def stride(x: Tensor):
-    return x.view(x.shape).stride()
+    return x.view(x.shape).stride()

blksprs/utils/validation.py

@@ -124,3 +124,7 @@ def _check_skip_validation():
 def _set_skip_validation(skip_validation: bool):
     global VALIDATION
     VALIDATION = not skip_validation
+
+
+def disable_validation():
+    _set_skip_validation(True)

{blksprs-1.8.2.dist-info → blksprs-1.8.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: blksprs
-Version: 1.8.2
+Version: 1.8.3
 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
@@ -22,6 +22,14 @@ Requires-Dist: matplotlib; extra == "test"
 [![GitHub Release](https://img.shields.io/github/v/release/FelixSchoen/blksprs?include_prereleases&label=Latest%20Release)](https://github.com/FelixSchoen/blksprs/releases)
 [![Python Version](https://img.shields.io/badge/Python%20Version-3.11-blue)](https://www.python.org/downloads/release/python-3119/)
 
+## Important Notice
+
+🚨 **Non-Final API** 🚨
+
+Although it already supports a wide variety of functions, this library is still under active development and the API is
+subject to change. For feature requests or bug reports, please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
+We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
+
 ## Overview
 
 A lightweight and efficient library for operations on block-sparse matrices in PyTorch using Triton.
@@ -51,14 +59,14 @@ These include, e.g.,
 Note that in order to correctly apply element-wise operations between two sparse tensors their sparsity layouts have to
 match.
 
-Further helpful operations (included in the ``bs.misc`` module) that do **not** support gradient calculation include:
+Further helpful operations (included in the ``bs.ops.misc`` module) that do **not** support gradient calculation include:
 
 - Row-wise sum, max, addition, and subtraction
 - Broadcast addition and subtraction between slices
 
 Furthermore, the library provides a set of utility functions for the creation of sparsity layouts based on existing
-dense tensors and for the scatter operation (module ``bs.layout``), as well as utility functions to ensure correct input
-dimensionality (module ``bs.util``).
+dense tensors and for the scatter operation (module ``bs.layouting``), as well as utility functions to apply linear layers,
+ensure correct input dimensionality, and validate input (module ``bs.utils``).
 
 ## Installation
 
@@ -111,14 +119,14 @@ def test_readme():
     y = torch.randn(size=(b, h, n, k), device="cuda").transpose(-1, -2).contiguous()
 
     # Convert tensors to three-dimensional (dense) tensors since Triton can only handle tensors of exactly three dimensions
-    x_dense, x_shape_original = bs.util.do_shape_blocksparse(x)
-    y_dense, y_shape_original = bs.util.do_shape_blocksparse(y)
+    x_dense, x_shape_original = bs.utils.do_shape_blocksparse(x)
+    y_dense, y_shape_original = bs.utils.do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    sparsity_layout_x = bs.layout.build_sparsity_layout(x_dense, sparsity_block_size,
-                                                        triton_block_size=triton_block_size)
-    sparsity_layout_y = bs.layout.build_sparsity_layout(y_dense, sparsity_block_size,
-                                                        triton_block_size=triton_block_size)
+    sparsity_layout_x = bs.layouting.build_sparsity_layout(x_dense, sparsity_block_size,
+                                                           triton_block_size=triton_block_size)
+    sparsity_layout_y = bs.layouting.build_sparsity_layout(y_dense, sparsity_block_size,
+                                                           triton_block_size=triton_block_size)
 
     # Create random sparsity layout for output tensor
     sparsity_layout_o = _get_random_sparsity_layout(b * h, m, n, sparsity_block_size, sparsity_percentage)
@@ -150,12 +158,12 @@ def test_readme():
     assert torch.allclose(o_dense, o_torch_round_trip, atol=2e-2)  # Note that small numerical differences are expected
 
     # Assert that the output has the correct sparsity layout
-    actual_sparsity_layout_o = bs.layout.build_sparsity_layout(o_dense, sparsity_block_size,
-                                                               triton_block_size=triton_block_size)
+    actual_sparsity_layout_o = bs.layouting.build_sparsity_layout(o_dense, sparsity_block_size,
+                                                                  triton_block_size=triton_block_size)
     assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
-    o = bs.util.undo_shape_blocksparse(o_dense, x_shape_original)
+    o = bs.utils.undo_shape_blocksparse(o_dense, x_shape_original)
 
     # Other available functions
     bs.transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)

blksprs-1.8.3.dist-info/RECORD

@@ -0,0 +1,23 @@
+blksprs/__init__.py,sha256=YMrERuEf1hTv5vVdOvPEzh9rESn4uqOB7WHB12Qs5lU,1836
+blksprs/layouting/distribution_layout.py,sha256=wmj1SwWyY_fhbvMmh6AXrR77LoSp6xLwUWCCyO9i5lk,4239
+blksprs/layouting/sparsity_layout.py,sha256=-sScIn4hhG35j9BXytrojEzp8jnFkMargJjtivPV1fc,9755
+blksprs/ops/conversion.py,sha256=ol-iV45wDzp9G1dJEkY53EdrvnmHzcl7QQmPJ-xqQTs,22410
+blksprs/ops/distribution.py,sha256=fXZV6UegCVpIwzh-A825OSYClHWu5k0UMYdO2UGDUpM,17067
+blksprs/ops/matmul.py,sha256=yh2ZnO0ZltT1AgadiFP0vX28YJ4n74xO-I_5vFUmOmA,11452
+blksprs/ops/partitioning.py,sha256=K0ExR2a3W62d_9xxCJzsdJDLgtbxTI6P8loOOBdhPzE,7674
+blksprs/ops/repeat.py,sha256=IvSIRbuyFn0b57LObymLgup0LqlWQ3ndIw-QuiYQcaU,14564
+blksprs/ops/softmax.py,sha256=CDQT2KnwkJ4hGIgT0EUp6P92uiYpCdJQ9zxcdgSAAJA,12102
+blksprs/ops/transpose.py,sha256=jxzFFffrj4S_9tiCrwwUMdz6EA98o1dziWXjlqb64a4,6859
+blksprs/ops/experimental/distribution_mdi.py,sha256=HaRUu6LTWATzjuHWgddIUE-0fgY-O87STpJO4JY7k_8,20357
+blksprs/ops/misc/broadcast_ops.py,sha256=cPtRJa3pkZfY1QG51CJ-zDn4SK-CRpX5LEXoKGGMvRU,5418
+blksprs/ops/misc/exp.py,sha256=FnSFosBfJHuiEbD0MD-i4axLghRn4a0f8KvHXrKBB6M,3802
+blksprs/ops/misc/row_wise.py,sha256=SvJuNww-_QoVKTyTjMvjmzHlBuUlTKamkuq_rKzwAqs,17081
+blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
+blksprs/utils/blksprs_tensor.py,sha256=VjplBgDhnf9sxf-1R5feA0xp5FDCDdaeZmCeoIRdCnc,151
+blksprs/utils/processing.py,sha256=hYsFxEbQKcbqU4WtZWusPnWMHg8ZAZF1SKZJYjez9aU,2060
+blksprs/utils/tools.py,sha256=r7Y4C37vfSWUyQTGwa8NyRqgovmsq9hMufkenqYHOxo,539
+blksprs/utils/validation.py,sha256=IZxH2HZpePmv7lRqLsSwV_6FwsdnTXv9q4j98vCMSsQ,4195
+blksprs-1.8.3.dist-info/METADATA,sha256=DZkJ_HeetF1V6-_F6GeG0uXT-QmttMFOq4ao8fiSMgQ,8458
+blksprs-1.8.3.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+blksprs-1.8.3.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-1.8.3.dist-info/RECORD,,

blksprs-1.8.2.dist-info/RECORD

@@ -1,22 +0,0 @@
-blksprs/__init__.py,sha256=np0msosWMaZNVVfuFGt8rE6HZURyIald391dKAs1dSQ,1093
-blksprs/experimental/distribution_mdi.py,sha256=HaRUu6LTWATzjuHWgddIUE-0fgY-O87STpJO4JY7k_8,20357
-blksprs/layouting/distribution_layout.py,sha256=wmj1SwWyY_fhbvMmh6AXrR77LoSp6xLwUWCCyO9i5lk,4239
-blksprs/layouting/sparsity_layout.py,sha256=-sScIn4hhG35j9BXytrojEzp8jnFkMargJjtivPV1fc,9755
-blksprs/misc/broadcast_ops.py,sha256=cPtRJa3pkZfY1QG51CJ-zDn4SK-CRpX5LEXoKGGMvRU,5418
-blksprs/misc/exp.py,sha256=FnSFosBfJHuiEbD0MD-i4axLghRn4a0f8KvHXrKBB6M,3802
-blksprs/misc/partitioning.py,sha256=K0ExR2a3W62d_9xxCJzsdJDLgtbxTI6P8loOOBdhPzE,7674
-blksprs/misc/row_wise.py,sha256=SvJuNww-_QoVKTyTjMvjmzHlBuUlTKamkuq_rKzwAqs,17081
-blksprs/ops/conversion.py,sha256=ol-iV45wDzp9G1dJEkY53EdrvnmHzcl7QQmPJ-xqQTs,22410
-blksprs/ops/distribution.py,sha256=fXZV6UegCVpIwzh-A825OSYClHWu5k0UMYdO2UGDUpM,17067
-blksprs/ops/matmul.py,sha256=yh2ZnO0ZltT1AgadiFP0vX28YJ4n74xO-I_5vFUmOmA,11452
-blksprs/ops/repeat.py,sha256=IvSIRbuyFn0b57LObymLgup0LqlWQ3ndIw-QuiYQcaU,14564
-blksprs/ops/softmax.py,sha256=D9wITz3KB24QXGGjgn_RLQ0Iiq_SjX0bTbUyv9479uU,12094
-blksprs/ops/transpose.py,sha256=jxzFFffrj4S_9tiCrwwUMdz6EA98o1dziWXjlqb64a4,6859
-blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
-blksprs/utils/blksprs_tensor.py,sha256=VjplBgDhnf9sxf-1R5feA0xp5FDCDdaeZmCeoIRdCnc,151
-blksprs/utils/tools.py,sha256=S3836Zuc-BMigv-5mLTjRznCzuaF6oYW-Ir9zzUnr3o,655
-blksprs/utils/validation.py,sha256=WzihRPibXYzss3PMkhDt5_d3Q3NHA_d1TzTz3CoGPGg,4136
-blksprs-1.8.2.dist-info/METADATA,sha256=Zoc860mYmFss7v5ChNoi9407v1qDo_ecc6JUWCvaesg,8009
-blksprs-1.8.2.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-blksprs-1.8.2.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.8.2.dist-info/RECORD,,