blksprs 0.2b4__py3-none-any.whl → 1.1__py3-none-any.whl

blksprs/utils/validation.py

@@ -0,0 +1,97 @@
+import torch
+from torch import Tensor
+
+
+def validate_dimensions(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
+    for tensor in tensors:
+        if tensor.dim() != 3:
+            raise ValueError("Tensor must have 3 dimensions")
+
+
+def validate_contiguous(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
+    for tensor in tensors:
+        if not tensor.is_contiguous():
+            raise ValueError("Tensor must be contiguous")
+
+
+def validate_dtype_float(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
+    for tensor in tensors:
+        if tensor.dtype != torch.float32:
+            raise ValueError("Tensor must have float32 dtype")
+
+
+def validate_dtype_int(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
+    for tensor in tensors:
+        if tensor.dtype != torch.int32 and tensor.dtype != torch.int64:
+            raise ValueError("Tensor must have int32 or int64 dtype")
+
+
+def validate_device(*tensors: Tensor) -> None:
+    if _skip_validation():
+        return
+
+    device = None
+
+    for i, tensor in enumerate(tensors):
+        if i == 0:
+            device = tensor.device
+
+            if not device.type == 'cuda':
+                raise ValueError("Tensors must be on GPU")
+
+        if tensor.device != device:
+            raise ValueError("Tensors must be on same device")
+
+
+def validate_sparsity(sparsity_block_size: int, *tensor_sparsity_layout_tuples: tuple[Tensor, Tensor]) -> None:
+    if _skip_validation():
+        return
+
+    for (tensor, sparsity_layout) in tensor_sparsity_layout_tuples:
+        _validate_sparsity_layout_values(sparsity_layout)
+
+        if not (tensor.size(-1) == tensor.size(-2) == sparsity_block_size):
+            raise ValueError("Blocks not conforming to sparsity block size")
+        if not tensor.size(0) == torch.sum(sparsity_layout.reshape(-1)):
+            raise ValueError("Mismatch between sparsity layout and blocks")
+
+
+def _validate_sparsity_layout_values(sparsity_layout: Tensor):
+    if not torch.all(torch.logical_or(sparsity_layout == 0, sparsity_layout == 1)):
+        raise ValueError("Sparsity layout values must be either 0 or 1")
+
+def validate_sparsity_block_size(sparsity_block_size: int, *tensors):
+    if _skip_validation():
+        return
+
+    if not (sparsity_block_size & (sparsity_block_size - 1)) == 0:
+        raise ValueError("Sparsity block size must be a power of 2")
+
+    for tensor in tensors:
+        if not (tensor.size(-1) % sparsity_block_size == 0 and tensor.size(-2) % sparsity_block_size == 0):
+            raise ValueError("Tensor sizes must be divisible by sparsity block size")
+
+def validate_triton_block_size(triton_block_size: int, sparsity_block_size: int):
+    if _skip_validation():
+        return
+
+    if triton_block_size is None:
+        return
+
+    if triton_block_size > sparsity_block_size:
+        raise ValueError("Triton block size cannot be larger than sparsity block size")
+
+def _skip_validation():
+    return False

blksprs-1.1.dist-info/METADATA

@@ -0,0 +1,164 @@
+Metadata-Version: 2.1
+Name: blksprs
+Version: 1.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
+Project-URL: Bugtracker, https://github.com/FelixSchoen/blksprs/issues
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Provides-Extra: deploy
+Requires-Dist: build; extra == "deploy"
+Requires-Dist: twine; extra == "deploy"
+Requires-Dist: pdoc3; extra == "deploy"
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-xdist; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: coverage; extra == "test"
+Requires-Dist: matplotlib; extra == "test"
+
+# blksprs
+
+## Overview
+
+A lightweight and efficient library for operations on block-sparse matrices in PyTorch using Triton.
+
+Currently supported operations (includes gradient calculation):
+
+- Sparse matrix multiplication (_supports any combination of sparse and dense matrices due to support
+  for `sparse = sparse @ sparse` matmul_)
+- Softmax
+- Transposition
+- Gather
+- Scatter (_supports either no reduction or summation, gradients are only available for summation_)
+- Conversion from and to sparse form
+
+As with this library sparse matrices are represented using a tuple of `(matrix, sparsity_layout, sparsity_block_size)`,
+any element-wise operations can be applied in regular torch-like fashion.
+These include, e.g.,
+
+- Element-wise addition and subtraction
+- Element-wise multiplication and division
+- Element-wise exponentiation
+- ...
+
+Note that in order to correctly apply element-wise operations between two sparse tensors their sparsity layouts have to
+match.
+
+Furthermore, the library provides a set of utility functions for the creation of sparsity layouts based on existing
+dense tensors.
+
+## Installation
+
+Note that due to the dependency on [Triton](https://github.com/triton-lang/triton) this library is only compatible with
+the Linux platform.
+
+We recommend installing blksprs from [PyPI](https://pypi.org/project/blksprs/) using pip:
+
+```pip install blksprs```
+
+## Changelog
+
+See [`CHANGELOG.md`](https://github.com/FelixSchoen/blksprs/blob/main/CHANGELOG.md) for a detailed changelog.
+
+## Usage
+
+We provide an example below to demonstrate the usage of the library.
+For more detailed examples, please refer to
+the [test cases](https://github.com/FelixSchoen/blksprs/blob/main/test/cases/test_blocksparse.py) which cover all
+implemented operations and functions.
+The example below can also be found in
+the [test cases](https://github.com/FelixSchoen/blksprs/blob/main/test/cases/test_readme.py).
+
+```python
+import torch
+
+from blksprs.layouting.sparsity_layout import build_sparsity_layout
+from blksprs.ops.conversion import to_sparse, to_dense
+from blksprs.ops.matmul import matmul
+from blksprs.ops.row_wise_sum import row_wise_sum
+from blksprs.ops.softmax import softmax
+from blksprs.ops.transpose import transpose
+from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse
+
+
+def test_readme():
+    # Set up parameters (batch size, number of heads, dimensions for matrices (m, k) and (n, k))
+    b, h, m, n, k = 2, 4, 64, 64, 16
+
+    # Percentage of blocks that will be sparse in the output for demonstration purposes
+    sparsity_percentage = 25
+
+    # Must be a power of two, greater than or equal to 16 for matmul, and divide m, n, and k
+    sparsity_block_size = 16
+
+    # Must be a power of two and smaller than or equal to sparsity_block_size
+    # If it is set to ``none`` a value will be chosen automatically
+    triton_block_size = None
+
+    # Initialise random (dense) tensors
+    x = torch.randn(size=(b, h, m, k), device="cuda")
+    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 = do_shape_blocksparse(x)
+    y_dense, y_shape_original = do_shape_blocksparse(y)
+
+    # Create sparsity layouts from existing tensors
+    sparsity_layout_x = build_sparsity_layout(x_dense, sparsity_block_size, triton_block_size=triton_block_size)
+    sparsity_layout_y = 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)
+
+    # Convert tensors to sparse tensors for matrix multiplication
+    x_sparse = to_sparse(x_dense, sparsity_layout_x, sparsity_block_size, triton_block_size=triton_block_size)
+    y_sparse = to_sparse(y_dense, sparsity_layout_y, sparsity_block_size, triton_block_size=triton_block_size)
+
+    # Perform matrix multiplication
+    o_sparse = matmul(x_sparse, sparsity_layout_x, y_sparse, sparsity_layout_y, sparsity_layout_o, sparsity_block_size,
+                      triton_block_size=triton_block_size)
+    o_dense = to_dense(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+
+    # Sanity check
+    o_torch = torch.matmul(x_dense, y_dense)
+
+    # Perform round trip to set sparse blocks to 0
+    o_torch_round_trip = to_dense(
+        to_sparse(o_torch, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size),
+        sparsity_layout_o, sparsity_block_size, fill_value=0, triton_block_size=triton_block_size)
+
+    # Assert that the output is correct
+    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 = build_sparsity_layout(o_dense, sparsity_block_size, triton_block_size=triton_block_size)
+    assert torch.allclose(actual_sparsity_layout_o, sparsity_layout_o)
+
+    # Convert output tensor back to original shape
+    o = undo_shape_blocksparse(o_dense, x_shape_original)
+
+    # Other available functions
+    transpose(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    softmax(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    row_wise_sum(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+
+
+def _get_random_sparsity_layout(b, m, n, sparsity_block_size, sparsity_percentage):
+    """Helper function, creates a random sparsity layout for a given shape with a given percentage of blocks marked as sparse.
+
+    """
+    m_s = m // sparsity_block_size
+    n_s = n // sparsity_block_size
+
+    sparsity_layout = torch.ones(size=(b, m_s, n_s), device="cuda", dtype=torch.int)
+
+    num_zero_elements = int(m_s * n_s * (sparsity_percentage / 100))
+    for b_i in range(b):
+        indices = torch.randperm(m_s * n_s)[:num_zero_elements]
+        sparsity_layout[b_i, indices // n_s, indices % n_s] = 0
+
+    return sparsity_layout
+```

blksprs-1.1.dist-info/RECORD

@@ -0,0 +1,17 @@
+blksprs/layouting/distribution_layout.py,sha256=GQ-ZRXbeImiLcbaqnL2FuUZ6DoFwmB0naT_YrOpD84Q,4940
+blksprs/layouting/sparsity_layout.py,sha256=Z9Ac88kQZVaPp27ymlLwyGN14ZfMyljXp6oM_gSsFMQ,2902
+blksprs/misc/broadcast_addition.py,sha256=vf1Hdqz9Uyqykto3DCjmdyepMzpMXL238SpANQqRAwI,5297
+blksprs/ops/conversion.py,sha256=COhHE5KvwhrtdUTLZX1wmxFe0kDNMY97iIhnkMmztBA,11362
+blksprs/ops/distribution.py,sha256=_fQb6fWpLxocAh86D74ATahChi0EK0eBb4eUOUEBVps,16769
+blksprs/ops/exp.py,sha256=qs8fVtCzxl4CKT4GepaqurjEL62jyi8VjMY12JFrFAU,3674
+blksprs/ops/matmul.py,sha256=x3lrYg4g8fIf5PeMtZY_SEpi11kP9RFcRoemCIxcSDE,11086
+blksprs/ops/row_wise_sum.py,sha256=ojuSejV37cLtRNS3lBfknA5KY3TEg8EHxOqVT6JZzoM,11387
+blksprs/ops/softmax.py,sha256=ZyeAVqmG_VzJ72FArGrpUSFfoSM4GPxyubrmNKERVIA,11654
+blksprs/ops/transpose.py,sha256=DVEXoxo2MoTNL3NZrjxsukMDrzk2vnEXL1uRnKFWkn0,6722
+blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
+blksprs/utils/tools.py,sha256=P2UALvccRjJJ7w05YGuaxB3qmNObgct4idfM0jlE2wg,465
+blksprs/utils/validation.py,sha256=gJYZO5C48YUrXV3Fy_Z_lCaOpiFj951FT-Od7sKfprg,3007
+blksprs-1.1.dist-info/METADATA,sha256=NIdEtqxj4SBUOP1eMlBz2RoOppwlQx9sJnRmDicWvp4,6982
+blksprs-1.1.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+blksprs-1.1.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-1.1.dist-info/RECORD,,

{blksprs-0.2b4.dist-info → blksprs-1.1.dist-info}/WHEEL

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