blksprs 1.11__py3-none-any.whl → 2.0__py3-none-any.whl

blksprs/ops/transpose.py

@@ -1,17 +1,16 @@
 import torch
-import triton
 from torch import Tensor
-from triton import language as tl
+from torch._library import triton_op
 
-from blksprs.ops.flow import flow_forward_pull
+from blksprs.ops.flow import flow_pull_forward
 from blksprs.utils.blksprs_tensor import BlksprsTensor
-from blksprs.utils.tools import get_triton_block_size, stride
 from blksprs.utils.validation import validate_dimensions, validate_contiguous, validate_device, \
-    validate_sparsity, validate_sparsity_block_size, validate_triton_block_size
+    validate_sparsity, validate_sparsity_block_size
 
 
-def transpose(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int, triton_block_size: int = None,
-              lut: dict = None) -> (BlksprsTensor, Tensor):
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
+def transpose(x: BlksprsTensor, sparsity_layout: Tensor,
+              sparsity_block_size: int, lut: dict = None) -> (BlksprsTensor, Tensor):
     """Transposes a block-sparse tensor in compressed form.
 
     Note:
@@ -21,7 +20,6 @@ def transpose(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: in
         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.
-        triton_block_size (int): The block size to use for the triton kernel (default ``None``).
         lut (dict, optional): A dictionary containing the look-up tables for the operation (default ``None``).
 
     Returns:
@@ -30,68 +28,73 @@ def transpose(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: in
 
     """
     x = x.contiguous()
-    x_t = x.transpose(-1, -2).contiguous()
 
     validate_dimensions(x)
     validate_contiguous(x)
     validate_device(x)
     validate_sparsity(sparsity_block_size, (x, sparsity_layout))
     validate_sparsity_block_size(sparsity_block_size, x)
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
 
-    lut = _BlocksparseTranspose.build_lut(lut, sparsity_layout)
+    lut = transpose_build_lut(lut, sparsity_layout)
 
-    return BlksprsTensor(
-        _BlocksparseTranspose.apply(x_t, lut["sparsity_layout_t"], lut["sparsity_lut"], lut["sparsity_reverse_lut"],
-                                    sparsity_block_size,
-                                    lut["n_sparse_blocks"], triton_block_size)), lut["sparsity_layout_t"]
+    return BlksprsTensor(transpose_forward(x, lut["sparsity_layout_t"],
+                                           lut["sparsity_lut"], lut["sparsity_reverse_lut"],
+                                           sparsity_block_size, lut["n_sparse_blocks"])), lut["sparsity_layout_t"]
 
 
-class _BlocksparseTranspose(torch.autograd.Function):
+@triton_op("blksprs::transpose_forward", mutates_args={})
+def transpose_forward(x: Tensor, sparsity_layout_o: Tensor,
+                      sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
+                      sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
+    with torch.no_grad():
+        x_t = x.transpose(-1, -2).contiguous()
+        return flow_pull_forward(x_t, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
+                                 sparsity_block_size, n_sparse_blocks)
 
-    @staticmethod
-    def build_lut(lut: dict, sparsity_layout: Tensor):
-        if lut is None:
-            lut = dict()
 
-        if "sparsity_layout_t" not in lut:
-            sparsity_layout_t = sparsity_layout.transpose(-1, -2).contiguous()
-            lut["sparsity_layout_t"] = sparsity_layout_t
+def transpose_wrapper_backward(ctx, grad_output):
+    sparsity_layout = ctx.saved_tensors[0]
+    sparsity_block_size = ctx.sparsity_block_size
 
-        if "sparsity_lut" not in lut:
-            sparsity_lut = torch.nonzero(lut["sparsity_layout_t"]).contiguous()
-            lut["sparsity_lut"] = sparsity_lut
+    return transpose(grad_output, sparsity_layout, sparsity_block_size)[
+        0], None, None, None, None, None
 
-        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()).transpose(-1, -2).contiguous().reshape(-1))
-            lut["sparsity_reverse_lut"] = sparsity_reverse_lut
 
-        if "n_sparse_blocks" not in lut:
-            n_sparse_blocks = torch.sum(sparsity_layout.to(torch.int)).item()
-            lut["n_sparse_blocks"] = n_sparse_blocks
+def transpose_build_lut(lut: dict, sparsity_layout: Tensor):
+    if lut is None:
+        lut = dict()
 
-        validate_contiguous(lut["sparsity_layout_t"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
+    if "sparsity_layout_t" not in lut:
+        sparsity_layout_t = sparsity_layout.transpose(-1, -2).contiguous()
+        lut["sparsity_layout_t"] = sparsity_layout_t
 
-        return lut
+    if "sparsity_lut" not in lut:
+        sparsity_lut = torch.nonzero(lut["sparsity_layout_t"]).contiguous()
+        lut["sparsity_lut"] = sparsity_lut
 
-    @staticmethod
-    def forward(ctx, x: Tensor, sparsity_layout_o: Tensor, sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
-                sparsity_block_size: int,
-                n_sparse_blocks: int, triton_block_size: int) -> Tensor:
-        ctx.save_for_backward(sparsity_layout_o)
+    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()).transpose(-1, -2).contiguous().reshape(-1))
+        lut["sparsity_reverse_lut"] = sparsity_reverse_lut
 
-        return flow_forward_pull(ctx, x, sparsity_layout_o, sparsity_lut, sparsity_reverse_lut,
-                                 sparsity_block_size, n_sparse_blocks, triton_block_size)
+    if "n_sparse_blocks" not in lut:
+        n_sparse_blocks = torch.sum(sparsity_layout.to(torch.int)).item()
+        lut["n_sparse_blocks"] = n_sparse_blocks
 
-    @staticmethod
-    def backward(ctx, grad_output):
-        sparsity_layout = ctx.saved_tensors[0]
-        sparsity_block_size = ctx.sparsity_block_size
-        triton_block_size = ctx.triton_block_size
+    validate_contiguous(lut["sparsity_layout_t"], lut["sparsity_lut"], lut["sparsity_reverse_lut"])
 
-        return transpose(grad_output, sparsity_layout, sparsity_block_size, triton_block_size)[
-            0], None, None, None, None, None, None
+    return lut
+
+
+# noinspection PyUnusedLocal
+def transpose_setup_context(ctx, inputs, output):
+    (_, sparsity_layout_o, _, _, sparsity_block_size, _) = inputs
+
+    ctx.save_for_backward(sparsity_layout_o)
+    ctx.sparsity_block_size = sparsity_block_size
+
+
+transpose_forward.register_autograd(transpose_wrapper_backward, setup_context=transpose_setup_context)

blksprs/utils/autotuning.py

@@ -0,0 +1,78 @@
+import os
+
+blksprs_autotune_mode = os.getenv("BLKSPRS_AUTOTUNE", "DEFAULT")
+
+if blksprs_autotune_mode == "DEFAULT":
+    autotune_parameters = [
+        (16, 3, 8),
+        (16, 4, 4),
+        (16, 5, 2),
+
+        (32, 3, 8),
+        (32, 4, 4),
+        (32, 5, 2),
+
+        (64, 3, 8),
+        (64, 4, 4),
+        (64, 5, 2),
+
+        (128, 3, 8),
+        (128, 4, 4),
+        (128, 5, 2),
+    ]
+elif blksprs_autotune_mode == "TEST":
+    autotune_parameters = [
+        (16, 3, 8),
+
+        (32, 3, 8),
+
+        (64, 3, 8),
+    ]
+else:
+    raise NotImplementedError(f"Unknown autotune mode: {blksprs_autotune_mode}")
+
+import torch
+import triton
+
+
+def prune_autotune_configs(autotune_configs, kernel_args, **kwargs):
+    sparsity_block_size = kernel_args["sparsity_block_size"]
+
+    pruned_configs = []
+
+    for config in autotune_configs:
+        if config.kwargs["TRITON_BLOCK_SIZE"] <= sparsity_block_size:
+            pruned_configs.append(config)
+
+    assert len(pruned_configs) > 0, f"No valid autotune configs found for sparsity block size {sparsity_block_size}"
+
+    return pruned_configs
+
+
+def prune_autotune_configs_conversion(autotune_configs, kernel_args, **kwargs):
+    sparsity_block_size_from = kernel_args["sparsity_block_size_from"]
+    sparsity_block_size_to = kernel_args["sparsity_block_size_to"]
+    sparsity_block_size = min(sparsity_block_size_from, sparsity_block_size_to)
+
+    pruned_configs = []
+
+    for config in autotune_configs:
+        if config.kwargs["TRITON_BLOCK_SIZE"] <= sparsity_block_size:
+            pruned_configs.append(config)
+
+    assert len(pruned_configs) > 0, f"No valid autotune configs found for sparsity block size {sparsity_block_size}"
+
+    return pruned_configs
+
+
+@torch.compile
+def get_autotune_configs():
+    global autotune_parameters
+
+    autotune_configs = []
+
+    for block_size, num_stages, num_warps in autotune_parameters:
+        autotune_configs.append(
+            triton.Config({"TRITON_BLOCK_SIZE": block_size}, num_stages=num_stages, num_warps=num_warps))
+
+    return autotune_configs

blksprs/utils/benchmarking.py

@@ -5,13 +5,13 @@ from matplotlib import pyplot as plt
 
 
 def benchmark(method_labels: list[str], func_input_generator: Callable,
-              matrix_sizes: list[int], sparsity_block_sizes: list[int], triton_block_sizes: list[int],
+              matrix_sizes: list[int], sparsity_block_sizes: list[int],
               *funcs_test_subject: Callable, y_lim_top: int = None):
     quantiles = [0.5, 0.2, 0.8]
     results = {}
 
-    for matrix_size, sparsity_block_size, triton_block_size in zip(matrix_sizes, sparsity_block_sizes, triton_block_sizes):
-        arguments = func_input_generator(matrix_size, sparsity_block_size, triton_block_size)
+    for matrix_size, sparsity_block_size in zip(matrix_sizes, sparsity_block_sizes):
+        arguments = func_input_generator(matrix_size, sparsity_block_size)
 
         for i, func_test_subject in enumerate(funcs_test_subject):
             func_ms_avg, func_ms_min, func_ms_max = triton.testing.do_bench(

blksprs/utils/processing.py

@@ -11,6 +11,7 @@ from blksprs.ops.repeat import repeat
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def apply_torch_linear(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block_size: int,
                        linear: nn.Linear, bias: nn.Parameter = None) -> (BlksprsTensor, Tensor):
     # Extract weight and bias
@@ -25,7 +26,7 @@ def apply_torch_linear(x: BlksprsTensor, sparsity_layout: Tensor, sparsity_block
 
     # 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)
+    xw = matmul(x, sparsity_layout, BlksprsTensor(w_t_bs.to(x.dtype)), sparsity_layout_w_t, sparsity_layout_xw, sparsity_block_size)
     interim = xw
 
     # Apply bias

blksprs/utils/tools.py

@@ -1,25 +1,24 @@
 import torch
 from torch import Tensor, Size
 
+# Capture scalar outputs for JIT compilation
+torch._dynamo.config.capture_scalar_outputs = True
 
-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
 
     return x.reshape((*shape[:-2], *x.shape[-2:]))
 
 
-def get_triton_block_size(sparsity_block_size: int, limit: int = 128):
-    return min(sparsity_block_size, limit)
-
-
 def stride(x: Tensor):
     if x.dim() == 2:
         return x.size(1), 1

blksprs/utils/validation.py

@@ -26,6 +26,23 @@ def validate_dtype_float(*tensors: Tensor) -> None:
     if _check_skip_validation():
         return
 
+    dtype = None
+
+    for i, tensor in enumerate(tensors):
+        if i == 0:
+            dtype = tensor.dtype
+
+        if tensor.dtype != torch.float16 and tensor.dtype != torch.float32:
+            raise ValueError("Tensor must have either float16 or float32 dtype")
+
+        if tensor.dtype != dtype:
+            raise ValueError("Tensors must have same dtype")
+
+
+def validate_dtype_float_32(*tensors: Tensor) -> None:
+    if _check_skip_validation():
+        return
+
     for tensor in tensors:
         if tensor.dtype != torch.float32:
             raise ValueError("Tensor must have float32 dtype")
@@ -38,7 +55,7 @@ def validate_dtype_int(*tensors: Tensor) -> None:
     for tensor in tensors:
         if (tensor.dtype !=
                 torch.int32 and tensor.dtype != torch.int64):
-            raise ValueError("Tensor must have int32 or int64 dtype")
+            raise ValueError("Tensor must have either int32 or int64 dtype")
 
 
 def validate_device(*tensors: Tensor) -> None:
@@ -51,7 +68,7 @@ def validate_device(*tensors: Tensor) -> None:
         if i == 0:
             device = tensor.device
 
-            if not device.type == 'cuda':
+            if not device.type == "cuda":
                 raise ValueError("Tensors must be on GPU")
 
         if tensor.device != device:
@@ -96,6 +113,9 @@ def validate_sparsity_block_size(sparsity_block_size: int, *tensors):
     if _check_skip_validation():
         return
 
+    if not sparsity_block_size >= 16:
+        raise ValueError("Sparsity block size must be at least 16")
+
     if not (sparsity_block_size & (sparsity_block_size - 1)) == 0:
         raise ValueError("Sparsity block size must be a power of 2")
 
@@ -104,20 +124,6 @@ def validate_sparsity_block_size(sparsity_block_size: int, *tensors):
             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 _check_skip_validation():
-        return
-
-    if triton_block_size is None:
-        return
-
-    if not (triton_block_size & (triton_block_size - 1)) == 0:
-        raise ValueError("Triton block size must be a power of 2")
-
-    if triton_block_size > sparsity_block_size:
-        raise ValueError("Triton block size cannot be larger than sparsity block size")
-
-
 def _check_skip_validation():
     return not VALIDATION
 

{blksprs-1.11.dist-info → blksprs-2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: blksprs
-Version: 1.11
+Version: 2.0
 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
@@ -14,9 +14,8 @@ Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-xdist; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
 Requires-Dist: coverage; extra == "test"
+Requires-Dist: build; extra == "test"
 Requires-Dist: matplotlib; extra == "test"
-Provides-Extra: build
-Requires-Dist: build; extra == "build"
 
 # blksprs
 
@@ -25,6 +24,13 @@ Requires-Dist: build; extra == "build"
 
 ## Overview
 
+### News
+
+🎉 ***Version 2.0 released***. blksprs now supports kernel auto-tuning, JIT compilation, specification of pre-calculated
+LUTs, autocasting, and makes use of `torch.library.triton_op()`!
+
+---
+
 A lightweight and efficient library for operations on block-sparse matrices in PyTorch using Triton.
 
 Currently supported operations (includes gradient calculation):
@@ -52,23 +58,25 @@ 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.ops.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 
+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.layouting``),
+  dense tensors and for the scatter operation (module ``bs.layouting``),
 - for the application of ``nn.Linear``, ``nn.Dropout``, and ``nn.LayerNorm`` layers to block-sparse tensors,
 - as well as utility functions to ensure correct input dimensionality, and validate input (module ``bs.utils``).
 
-_* see the [Roadmap](#roadmap) section for more information_ 
+_* see the [Roadmap](#roadmap) section for more information_
 
 ## Installation
 
-Note that due to the dependency on [Triton](https://github.com/triton-lang/triton) this library is **only compatible with
+Note that due to the dependency on [Triton](https://github.com/triton-lang/triton) this library is **only compatible
+with
 the Linux platform**.
 Keep track of this [issue](https://github.com/triton-lang/triton/issues/1640) for updates.
 
@@ -78,8 +86,8 @@ We recommend installing blksprs from [PyPI](https://pypi.org/project/blksprs/) u
 
 ### Dependencies
 
-- [PyTorch](https://pytorch.org/) (built with v2.5.1)
-- _[NumPy](https://numpy.org/) (to get rid of warnings, built with v2.2.0)_
+- [PyTorch](https://pytorch.org/) (built with v2.6)
+- _[NumPy](https://numpy.org/) (to get rid of warnings, built with v2.2.4)_
 - _[Triton](https://github.com/triton-lang/triton) (included with PyTorch)_
 
 ## Changelog
@@ -89,12 +97,27 @@ See [`CHANGELOG.md`](https://github.com/FelixSchoen/blksprs/blob/main/CHANGELOG.
 ## Roadmap
 
 Note that since this library covers all our current needs it is in a **bugfix-only** state.
-This means that there are no plans to add new features, e.g., support for dimension specification of the ``split`` and ``merge`` operations.
+This means that there are no plans to add new features, e.g., support for dimension specification of the ``split`` and
+``merge`` operations.
 We will continue to maintain the library and fix any issues that arise.
 Should you find any bugs please open an [issue](https://github.com/FelixSchoen/blksprs/issues).
 We also encourage [pull requests](https://github.com/FelixSchoen/blksprs/pulls).
 
-It might be that this changes with future projects, but as of December 2024, we are content with the current state of the library.
+It might be that this changes with future projects, but as of March 2025, we are content with the current state of the
+library.
+
+## Known Limitations and Issues
+
+- Triton has a bug with `tl.atomix_max()` used for the row-wise max operation.
+  In order to work around this bug a manual conversion of some values is needed, (slightly) negatively impacting
+  performance.
+  Watch the [issue](https://github.com/triton-lang/triton/issues/6376) on Triton's issue tracker for more information.
+- PyTorch's `wrap_triton()` currently does not support config pruning. It thus cannot be used for some of the kernels,
+  which could impact graph compilation.
+- There seem to be some issues with autocasting, forcing some operations to manually cast.
+- There will be some slight numerical differences between vanilla and blksprs operations.
+  These instabilities are due to Triton and thus cannot be fixed by this library alone.
+  However, for all intents and purposes, these very minor differences should not matter and can safely be ignored.
 
 ## Usage
 
@@ -120,10 +143,6 @@ def test_readme():
     # 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()
@@ -133,53 +152,53 @@ def test_readme():
     y_dense, y_shape_original = bs.utils.do_shape_blocksparse(y)
 
     # Create sparsity layouts from existing tensors
-    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)
+    sparsity_layout_x = bs.layouting.build_sparsity_layout(x_dense, sparsity_block_size)
+    sparsity_layout_y = bs.layouting.build_sparsity_layout(y_dense, sparsity_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 = bs.to_sparse(x_dense, sparsity_layout_x, sparsity_block_size, triton_block_size=triton_block_size)
-    y_sparse = bs.to_sparse(y_dense, sparsity_layout_y, sparsity_block_size, triton_block_size=triton_block_size)
+    x_sparse = bs.ops.to_sparse(x_dense, sparsity_layout_x, sparsity_block_size)
+    y_sparse = bs.ops.to_sparse(y_dense, sparsity_layout_y, sparsity_block_size)
+
+    # As of version 2.0, blksprs supports JIT compilation
+    matmul_compiled = torch.compile(bs.ops.matmul)
 
     # Perform matrix multiplication
-    o_sparse = bs.matmul(x_sparse, sparsity_layout_x, y_sparse, sparsity_layout_y, sparsity_layout_o,
-                         sparsity_block_size,
-                         triton_block_size=triton_block_size)
+    o_sparse = matmul_compiled(x_sparse, sparsity_layout_x,
+                               y_sparse, sparsity_layout_y,
+                               sparsity_layout_o, sparsity_block_size)
 
     # Apply element-wise operation
     o_sparse = torch.add(o_sparse, 1)
 
-    o_dense = bs.to_dense(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    o_dense = bs.ops.to_dense(o_sparse, sparsity_layout_o, sparsity_block_size)
 
     # Sanity check
     o_torch = torch.matmul(x_dense, y_dense)
     o_torch = torch.add(o_torch, 1)
 
     # Perform round trip to set sparse blocks to 0
-    o_torch_round_trip = bs.to_dense(
-        bs.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)
+    o_torch_round_trip = bs.ops.to_dense(
+        bs.ops.to_sparse(o_torch, sparsity_layout_o, sparsity_block_size),
+        sparsity_layout_o, sparsity_block_size, fill_value=0)
 
     # 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 = bs.layouting.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)
     assert torch.allclose(actual_sparsity_layout_o.to(torch.int), sparsity_layout_o)
 
     # Convert output tensor back to original shape
     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)
-    bs.softmax(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
-    bs.misc.row_wise_sum(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
-    bs.misc.row_wise_max(o_sparse, sparsity_layout_o, sparsity_block_size, triton_block_size=triton_block_size)
+    bs.ops.transpose(o_sparse, sparsity_layout_o, sparsity_block_size)
+    bs.ops.softmax(o_sparse, sparsity_layout_o, sparsity_block_size)
+    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)
 
 
 def _get_random_sparsity_layout(b, m, n, sparsity_block_size, sparsity_percentage):

blksprs-2.0.dist-info/RECORD

@@ -0,0 +1,23 @@
+blksprs/__init__.py,sha256=QXYiDairqn5LzvpoNyBXHdQfyac-aX3IDjTQUMkaH-A,1598
+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/distribution.py,sha256=0tPldv0ARzmCV1CU2jvfqpHBgOuHPrDFiCtqsLs7CZc,20789
+blksprs/ops/flow.py,sha256=PDZAD8u4y9qW1IXERki6ItKbEKnm_ChG8SKWM3_P9Oc,8245
+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/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
+blksprs/utils/autotuning.py,sha256=a-kmWRjJ3eED2XbjkQeOJSyW8bdIs27HgKMPvAKqWeU,2052
+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/validation.py,sha256=G8eQlvJVMKfEX3k2AwBD0A6Ck-gFoRLpLNY6HXsB3fA,4348
+blksprs-2.0.dist-info/METADATA,sha256=7YVN_akf-ewrAW5thDZzhT3hogn0dVxV73lvlPMk59c,9506
+blksprs-2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+blksprs-2.0.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
+blksprs-2.0.dist-info/RECORD,,

{blksprs-1.11.dist-info → blksprs-2.0.dist-info}/WHEEL

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

blksprs/utils/layout_utils.py

@@ -1,17 +0,0 @@
-import math
-
-import torch
-import triton
-from torch import Tensor
-from torch.xpu import device
-from triton import language as tl
-
-from blksprs.utils.blksprs_tensor import BlksprsTensor
-from blksprs.utils.tools import get_triton_block_size, stride
-from blksprs.utils.validation import validate_triton_block_size, validate_dimensions, validate_device, \
-    validate_contiguous, validate_sparsity, validate_sparsity_block_size
-
-
-def build_full_sparsity_layout(x: Tensor, sparsity_block_size: int) -> Tensor:
-    return torch.ones(size=(x.size(0), x.size(1) // sparsity_block_size, x.size(2) // sparsity_block_size),
-                      dtype=torch.bool, device=x.device)

blksprs-1.11.dist-info/RECORD

@@ -1,23 +0,0 @@
-blksprs/__init__.py,sha256=AJYVfR40nOfE5F3waHPVSuajwYDcoGkiEQc8HhQbUBU,1721
-blksprs/layouting/distribution_layout.py,sha256=xDGY5-J7uSD8oenlf8bEJ2amMiQG3NBf2klTTydbTJE,5140
-blksprs/layouting/sparsity_layout.py,sha256=IVtHc_nN3ZM2y4GFcys70PqDWmWc7tkHlVGlToErANk,9894
-blksprs/ops/conversion.py,sha256=QFtZ-nmY2JAWutheiO07vatXqz3eSZBP5Ym_U2Q1oWk,23299
-blksprs/ops/distribution.py,sha256=nHTuE7Tq0Q404VN8bWNC2sEwmmdAtgZI6I7auRICdps,21749
-blksprs/ops/flow.py,sha256=7tOXfTBKOAixYmDa_VXg7TwviLV5ZQMHQjtbyOjqA00,7879
-blksprs/ops/matmul.py,sha256=eVj_BGj78bJkXYuvw4KctMfcfveQBt5OdYmeXzdpO88,12631
-blksprs/ops/partitioning.py,sha256=qMv9w3yFWXwXIhIppdcJ_JMsoZ25HCH38vb6GRneoLM,10416
-blksprs/ops/repeat.py,sha256=i824ijprfYpCaEjiSG5FTUZz7wMS5ksVy_-vY7ZX8Fg,9729
-blksprs/ops/softmax.py,sha256=_mGkA2jHN8cXwtWXYswobEPyM7UC0JyzRszoE4ZYs7w,13063
-blksprs/ops/transpose.py,sha256=O1XhGIGiVkhOSKcBD0HrYaeK6HmpvEEzLb7zJl7xsIM,4246
-blksprs/ops/misc/broadcast_ops.py,sha256=pv0nssSDOdDbQFttpqUIs2ZXShqfm2RYCfJH-C5x3H0,5544
-blksprs/ops/misc/row_wise.py,sha256=DnV5-xEJUbqZlK2fETwHiPQDUMwT-lkc0VUhBlnJ5Y0,17458
-blksprs/utils/benchmarking.py,sha256=4pLVlnPW_2EM-NT3n4SClaRznVYEljztLbJcccz8kZE,1360
-blksprs/utils/blksprs_tensor.py,sha256=pfoz59aJixj_fIoFx76ySiygwRQUemmgjMKepZ2c4j0,244
-blksprs/utils/layout_utils.py,sha256=49ZdPS_gMn_IrWty3FARbi2rda5a8g5DmAEL8LOrC30,670
-blksprs/utils/processing.py,sha256=WLuMJQ8v-YovXwcDjhlDn3N31WMZXrtyeeyKSgq_zn4,3642
-blksprs/utils/tools.py,sha256=k2OfEplbQiAwVjP84zZf7SNB8FzvMtOFBL9sC98OCbI,683
-blksprs/utils/validation.py,sha256=CbxBbeQWJo8wox5eMoVzaTlP9FVBwt3-gxUOmi3EUgw,4213
-blksprs-1.11.dist-info/METADATA,sha256=NUEiHexWiFNbMxQI2TUEzMw9iGBhxqflhWr2xCgOw28,9105
-blksprs-1.11.dist-info/WHEEL,sha256=52BFRY2Up02UkjOa29eZOS2VxUrpPORXg1pkohGGUS8,91
-blksprs-1.11.dist-info/top_level.txt,sha256=qyp0IHeY3H2GQA97i4hk_To5rRBS2YcE1HRPSLy04fk,8
-blksprs-1.11.dist-info/RECORD,,