blksprs 1.10.2__tar.gz → 2.0__tar.gz

{blksprs-1.10.2 → blksprs-2.0}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: blksprs
-Version: 1.10.2
+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-1.10.2 → blksprs-2.0}/README.md

@@ -5,6 +5,13 @@
 
 ## 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):
@@ -32,23 +39,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.
 
@@ -58,8 +67,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
@@ -69,12 +78,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
 
@@ -100,10 +124,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()
@@ -113,53 +133,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-1.10.2 → blksprs-2.0}/blksprs/__init__.py

@@ -1,5 +1,7 @@
 from blksprs.utils.blksprs_tensor import BlksprsTensor
 
+__version__ = "2.0"
+
 
 class ops:
     from blksprs.ops.conversion import to_dense, to_sparse, from_blksprs, to_blksprs, adapt_layout
@@ -13,25 +15,21 @@ class ops:
     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 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
-    from blksprs.utils.layout_utils import build_full_sparsity_layout
+        build_sparsity_layout_matmul, build_sparsity_layout_matmul_fast, build_sparsity_layout_full
 
 
 class utils:
     from blksprs.utils.processing import apply_torch_linear, apply_torch_normalisation, apply_torch_dropout, \
         apply_function_applicable_row_wise
     from blksprs.utils.tools import do_shape_blocksparse, undo_shape_blocksparse
-    from blksprs.utils.validation import disable_validation
 
     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
+            validate_sparsity_block_size

{blksprs-1.10.2 → blksprs-2.0}/blksprs/layouting/distribution_layout.py

@@ -1,17 +1,23 @@
+import typing
+
 import torch
 import triton
 from torch import Tensor
+from torch._library import triton_op
+from torch._library.triton import wrap_triton
 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, \
+from blksprs.utils.tools import stride
+from blksprs.utils.autotuning import get_autotune_configs, prune_autotune_configs
+from blksprs.utils.validation import validate_dimensions, validate_device, \
     validate_contiguous
 
 
+@torch.amp.custom_fwd(device_type="cuda", cast_inputs=torch.float16)
 def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: Tensor,
                               dim: int, size_target: torch.Size,
-                              sparsity_block_size: int, triton_block_size: int = None) -> Tensor:
+                              sparsity_block_size: int) -> Tensor:
     """Builds the sparsity layout of either the source of a gather or the target of a scatter operation.
 
     Args:
@@ -20,7 +26,6 @@ def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: T
         dim (int): The dimension along which the operation is conducted.
         size_target (torch.Size): The size of the block-sparse target tensor in regular form.
         sparsity_block_size (int): The size of the sparsity blocks.
-        triton_block_size (int, optional): The block size to use for the triton kernel (default ``None``).
 
     Returns:
         Tensor: The sparsity layout of the source or target tensor.
@@ -34,49 +39,58 @@ def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: T
 
     adjusted_dim = dim % 3
 
-    output = torch.zeros(size_target[0], size_target[1] // sparsity_block_size, size_target[2] // sparsity_block_size,
-                         dtype=torch.bool, device=indices.device)
-
-    i_b, i_r, i_c = indices.size()
-    i_b_s, i_r_s, i_c_s = stride(indices)
-    s_lut_i_r, s_lut_i_c = sparsity_lut_i.size()
-    s_lut_i_r_s, s_lut_i_c_s = stride(sparsity_lut_i)
-    o_b, o_r, o_c = output.size()
-    o_b_s, o_r_s, o_c_s = stride(output)
-
-    if triton_block_size is None:
-        triton_block_size = get_triton_block_size(sparsity_block_size)
-
-    validate_triton_block_size(triton_block_size, sparsity_block_size)
-
-    triton_grid = lambda meta: [i_b,
-                                triton.cdiv(i_r, meta["TRITON_BLOCK_SIZE"]),
-                                triton.cdiv(i_c, meta["TRITON_BLOCK_SIZE"])]
-
-    (kernel_distribution_layout[triton_grid]
-     (indices,
-      i_b, i_b_s, i_r_s, i_c_s,
-      sparsity_lut_i,
-      s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
-      adjusted_dim,
-      output,
-      o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size,
-      triton_block_size))
-
-    return output
-
-
+    return build_distribution_layout_operation(indices, sparsity_lut_i, adjusted_dim, size_target, sparsity_block_size)
+
+
+@triton_op("blksprs::build_distribution_layout", mutates_args={})
+def build_distribution_layout_operation(indices: Tensor, sparsity_lut_i: Tensor,
+                                        adjusted_dim: int, size_target: typing.List[int],
+                                        sparsity_block_size: int) -> Tensor:
+    with torch.no_grad():
+        output = torch.zeros(size_target[0], size_target[1] // sparsity_block_size,
+                             size_target[2] // sparsity_block_size,
+                             dtype=torch.bool, device=indices.device)
+
+        i_b, i_r, i_c = indices.size()
+        i_b_s, i_r_s, i_c_s = stride(indices)
+        s_lut_i_r, s_lut_i_c = sparsity_lut_i.size()
+        s_lut_i_r_s, s_lut_i_c_s = stride(sparsity_lut_i)
+        o_b, o_r, o_c = output.size()
+        o_b_s, o_r_s, o_c_s = stride(output)
+
+        triton_grid = lambda meta: [i_b,
+                                    triton.cdiv(i_r, meta["TRITON_BLOCK_SIZE"]),
+                                    triton.cdiv(i_c, meta["TRITON_BLOCK_SIZE"])]
+
+        (wrap_triton(build_distribution_layout_kernel)[triton_grid]
+         (indices,
+          i_b, i_b_s, i_r_s, i_c_s,
+          sparsity_lut_i,
+          s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
+          adjusted_dim,
+          output,
+          o_b, o_b_s, o_r_s, o_c_s,
+          sparsity_block_size))
+
+        return output
+
+
+@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 kernel_distribution_layout(i,
-                               i_b, i_b_s, i_r_s, i_c_s,
-                               s_lut_i,
-                               s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
-                               dim,
-                               o,
-                               o_b, o_b_s, o_r_s, o_c_s,
-                               sparsity_block_size,
-                               TRITON_BLOCK_SIZE: tl.constexpr) -> None:
+def build_distribution_layout_kernel(i,
+                                     i_b, i_b_s, i_r_s, i_c_s,
+                                     s_lut_i,
+                                     s_lut_i_r, s_lut_i_r_s, s_lut_i_c_s,
+                                     dim,
+                                     o,
+                                     o_b, o_b_s, o_r_s, o_c_s,
+                                     sparsity_block_size,
+                                     TRITON_BLOCK_SIZE: tl.constexpr) -> None:
     # Get triton block indices
     pid_blk = tl.program_id(axis=0)
     pid_row = tl.program_id(axis=1)
@@ -98,7 +112,8 @@ def kernel_distribution_layout(i,
     blk_i_idx = (pid_blk * i_b_s +
                  ((pid_row * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_r_s)[:, None] +
                  ((pid_col * TRITON_BLOCK_SIZE + tl.arange(0, TRITON_BLOCK_SIZE)) * i_c_s)[None, :])
-    blk_i_msk = (blk_i_idx >= 0 and blk_i_idx < i_b * i_b_s)
+    blk_i_msk = (blk_i_idx >= 0 and
+                 blk_i_idx < i_b * i_b_s)
     blk_i = tl.load(i + blk_i_idx, mask=blk_i_msk)
 
     dst_bat_idx = tl.full((TRITON_BLOCK_SIZE, TRITON_BLOCK_SIZE), spa_bat_i, dtype=tl.int32)
@@ -116,5 +131,6 @@ def kernel_distribution_layout(i,
     blk_o_idx = ((dst_bat_idx * o_b_s) +
                  (dst_row_idx * o_r_s) +
                  (dst_col_idx * o_c_s))
-    blk_o_msk = (blk_o_idx >= 0 and blk_o_idx < o_b * o_b_s)
+    blk_o_msk = (blk_o_idx >= 0 and
+                 blk_o_idx < o_b * o_b_s)
     tl.store(o + blk_o_idx, blk_v, mask=blk_o_msk)