blksprs 1.11__tar.gz → 2.0rc2__tar.gz

{blksprs-1.11 → blksprs-2.0rc2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: blksprs
-Version: 1.11
+Version: 2.0rc2
 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, 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,14 @@ 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.
 
 ## Usage
 
@@ -120,10 +130,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 +139,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.11 → blksprs-2.0rc2}/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, 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,14 @@ 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.
 
 ## Usage
 
@@ -100,10 +111,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 +120,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.11 → blksprs-2.0rc2}/blksprs/__init__.py

@@ -18,19 +18,16 @@ class ops:
 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.11 → blksprs-2.0rc2}/blksprs/layouting/distribution_layout.py

@@ -4,14 +4,14 @@ from torch import Tensor
 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, get_autotune_configs
+from blksprs.utils.validation import validate_dimensions, validate_device, \
     validate_contiguous
 
 
 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 +20,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.
@@ -44,16 +43,11 @@ def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: T
     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]
+    (build_distribution_layout_kernel[triton_grid]
      (indices,
       i_b, i_b_s, i_r_s, i_c_s,
       sparsity_lut_i,
@@ -61,27 +55,34 @@ def build_distribution_layout(indices: BlksprsTensor, sparsity_layout_indices: T
       adjusted_dim,
       output,
       o_b, o_b_s, o_r_s, o_c_s,
-      sparsity_block_size,
-      triton_block_size))
+      sparsity_block_size))
 
     return output
 
 
+@triton.autotune(
+    configs=get_autotune_configs(),
+    key=[],
+    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)
     pid_col = tl.program_id(axis=2)
 
+    # Get valid triton block size
+    val_tbs = min(sparsity_block_size, TRITON_BLOCK_SIZE)
+
     # Get position of current sparsity block consisting of its batch, row, and column index
     spa_bat_i_idx = (pid_blk * s_lut_i_r_s + 0 * s_lut_i_c_s)
     spa_bat_i_msk = (spa_bat_i_idx >= 0 and spa_bat_i_idx < s_lut_i_r * s_lut_i_r_s)
@@ -96,9 +97,12 @@ def kernel_distribution_layout(i,
     spa_col_i = tl.load(s_lut_i + spa_col_i_idx, mask=spa_col_i_msk)
 
     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)
+                 ((pid_row * val_tbs + tl.arange(0, TRITON_BLOCK_SIZE)) * i_r_s)[:, None] +
+                 ((pid_col * val_tbs + 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) and
+                 (tl.arange(0, TRITON_BLOCK_SIZE)[:, None] < val_tbs and
+                  tl.arange(0, TRITON_BLOCK_SIZE)[None, :] < val_tbs))
     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 +120,8 @@ 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) and
+                 (tl.arange(0, TRITON_BLOCK_SIZE)[:, None] < val_tbs and
+                  tl.arange(0, TRITON_BLOCK_SIZE)[None, :] < val_tbs))
     tl.store(o + blk_o_idx, blk_v, mask=blk_o_msk)