blksprs 2.0rc3__tar.gz → 2.0rc4__tar.gz

{blksprs-2.0rc3 → blksprs-2.0rc4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.0rc3
+Version: 2.0rc4
 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

{blksprs-2.0rc3 → blksprs-2.0rc4}/blksprs/ops/conversion.py

@@ -54,7 +54,7 @@ def to_sparse(x: Tensor, sparsity_layout: Tensor,
 @triton_op("blksprs::to_sparse", mutates_args={})
 def to_sparse_forward(x: Tensor, _: Tensor,
                       sparsity_lut: Tensor, sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
-    output = torch.empty(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
+    output = torch.zeros(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
                          dtype=x.dtype, device=x.device)
 
     x_b, x_r, x_c = x.size()
@@ -87,6 +87,7 @@ def to_sparse_backward(ctx, grad_output):
 @triton.autotune(
     configs=get_autotune_configs(),
     key=[],
+    reset_to_zero=["o"]
 )
 @triton.jit
 def to_sparse_kernel(x,
@@ -252,6 +253,7 @@ def to_dense_backward(ctx, grad_output):
 @triton.autotune(
     configs=get_autotune_configs(),
     key=[],
+    restore_value=["o"]
 )
 @triton.jit
 def to_dense_kernel(x,

{blksprs-2.0rc3 → blksprs-2.0rc4}/blksprs/ops/distribution.py

@@ -54,7 +54,7 @@ def gather(src: BlksprsTensor, sparsity_layout_src: Tensor,
 def gather_forward(x: Tensor, sparsity_layout_x: Tensor, sparsity_reverse_lut_x: Tensor,
                    dim: int, i: Tensor, _: Tensor, sparsity_lut_i: Tensor,
                    sparsity_block_size: int) -> Tensor:
-    output = torch.empty_like(i, dtype=x.dtype)
+    output = torch.zeros_like(i, dtype=x.dtype)
 
     x_b, x_r, x_c = x.size()
     x_b_s, x_r_s, x_c_s = stride(x)
@@ -101,6 +101,7 @@ def gather_backward(ctx, grad_output):
 @triton.autotune(
     configs=get_autotune_configs(),
     key=[],
+    reset_to_zero=["o"]
 )
 @triton.jit
 def gather_kernel(x,

{blksprs-2.0rc3 → blksprs-2.0rc4}/blksprs/ops/flow.py

@@ -12,7 +12,7 @@ from blksprs.utils.tools import stride, get_autotune_configs
 def flow_pull_forward(x: Tensor, sparsity_layout_o: Tensor,
                       sparsity_lut: Tensor, sparsity_reverse_lut: Tensor,
                       sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
-    output = torch.empty(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
+    output = torch.zeros(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
                          dtype=x.dtype, device=x.device)
 
     x_b, x_r, x_c = x.size()
@@ -44,6 +44,7 @@ def flow_pull_forward(x: Tensor, sparsity_layout_o: Tensor,
 @triton.autotune(
     configs=get_autotune_configs(),
     key=[],
+    reset_to_zero=["o"]
 )
 @triton.jit
 def flow_pull_kernel(x,

{blksprs-2.0rc3 → blksprs-2.0rc4}/blksprs/ops/matmul.py

@@ -60,7 +60,7 @@ def matmul_forward(x: Tensor, y: Tensor,
                    sparsity_layout_y: Tensor, sparsity_reverse_lut_y: Tensor,
                    _: Tensor, sparsity_lut_o: Tensor,
                    sparsity_block_size: int, n_sparse_blocks: int) -> Tensor:
-    output = torch.empty(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
+    output = torch.zeros(size=(n_sparse_blocks, sparsity_block_size, sparsity_block_size),
                          dtype=x.dtype, device=x.device)
 
     x_b, x_r, x_c = x.size()
@@ -118,6 +118,7 @@ def matmul_backward(ctx, grad_output):
 @triton.autotune(
     configs=get_autotune_configs(),
     key=[],
+    reset_to_zero=["o"]
 )
 @triton.jit
 def matmul_kernel(x,

{blksprs-2.0rc3 → blksprs-2.0rc4}/blksprs/ops/misc/row_wise.py

@@ -354,7 +354,7 @@ def row_wise_sub(x: BlksprsTensor, sparsity_layout_x: Tensor, y: Tensor,
 def row_wise_add_forward(x: Tensor, sparsity_lut_x: Tensor,
                          sparsity_layout_x_rwm: Tensor, sparsity_reverse_x_lut_rwm: Tensor,
                          y: Tensor, sparsity_block_size: int) -> Tensor:
-    output = torch.empty_like(x)
+    output = torch.zeros_like(x)
 
     x_b, x_r, x_c = x.size()
     x_b_s, x_r_s, x_c_s = stride(x)
@@ -387,7 +387,8 @@ def row_wise_add_forward(x: Tensor, sparsity_lut_x: Tensor,
 
 @triton.autotune(
     configs=get_autotune_configs(),
-    key=[]
+    key=[],
+    reset_to_zero=["o"]
 )
 @triton.jit
 def kernel_blocksparse_row_wise_add(x,

{blksprs-2.0rc3 → blksprs-2.0rc4}/blksprs/ops/softmax.py

@@ -51,7 +51,7 @@ def softmax_forward(x: Tensor, sparsity_layout: Tensor,
                     sparsity_lut: Tensor,
                     sparsity_reverse_lut_rws: Tensor,
                     sparsity_block_size: int) -> Tensor:
-    output = torch.empty_like(x)
+    output = torch.zeros_like(x)
 
     x_b, x_r, x_c = x.size()
     x_b_s, x_r_s, x_c_s = stride(x)
@@ -108,7 +108,7 @@ def softmax_backward(ctx, grad_output):
     s_l_s_b, s_l_s_r, s_l_s_c = sparsity_layout_s.size()
     s_l_s_b_s, s_l_s_r_s, s_l_s_c_s = stride(sparsity_layout_s)
 
-    grad_x = torch.empty_like(o, dtype=torch.float)
+    grad_x = torch.zeros_like(o, dtype=torch.float)
 
     triton_grid = lambda meta: [o_b,
                                 triton.cdiv(o_r, meta["TRITON_BLOCK_SIZE"]),
@@ -133,7 +133,8 @@ def softmax_backward(ctx, grad_output):
 
 @triton.autotune(
     configs=get_autotune_configs(),
-    key=[]
+    key=[],
+    reset_to_zero=["o"]
 )
 @triton.jit
 def softmax_kernel(x,
@@ -198,7 +199,8 @@ def softmax_kernel(x,
 
 @triton.autotune(
     configs=get_autotune_configs(),
-    key=[]
+    key=[],
+    reset_to_zero=["o"]
 )
 @triton.jit
 def softmax_kernel_grad(g,

{blksprs-2.0rc3 → blksprs-2.0rc4}/blksprs.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: blksprs
-Version: 2.0rc3
+Version: 2.0rc4
 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

{blksprs-2.0rc3 → blksprs-2.0rc4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "blksprs"
-version = "2.0-rc.3"
+version = "2.0-rc.4"
 authors = [{ name = "Felix Schön", email = "schoen@kr.tuwien.ac.at" }]
 description = "A lightweight library for operations on blocksparse matrices in PyTorch."
 readme = "README.md"