question on blocksparse. about triton HOT 10 CLOSED

Sure! You can find a block-sparse attention example there https://github.com/ptillet/triton/blob/master/python/test/test_blocksparse.py#L150-L159. You can create a block-sparse softmax operation as follows:

sparse_softmax = triton.ops.blocksparse.softmax(layout, block)

and then call it on the output of a SDD matmul

from triton.

This is general block-sparse spMM. There are three modes:

• SDD: sparse = dense x dense, a.k.a. sampled dense-dense matrix multiplication
• DSD: dense = sparse x dense, the lhs is sparse
• DDS: dense = dense x sparse, the rhs is sparse

The output of SDD is in a block-sparse format that can be re-used for triton.ops.blocksparse.softmax and also by DSD for attention mechanisms.

from triton.

thx

from triton.

Hi what's the definition of block here? Is it to blockify the matrix and then do the MM?
I'm looking to implement the block sparse attention proposed by bigbird using triton ops. I want to know how to obtain the block in triton. Thanks

from triton.

The sparsity layout is specified as a tensor of 0s and 1s. On your example, this would only work if each colored square corresponds to a 16x16 (or 32x32, 64x64, 128x128) block of data. Also note that triton.ops.blocksparse doesn't support overlapping blocks

from triton.

Thanks. There won't be any overlapping blocks during attention. And the block size in default is 64X64.
Does triton provide any optimized softmax or layernorm kernel for the output from block sparse MM?

from triton.

No layernorm, but you can use triton.ops.blocksparse.softmax to reduce the row of the output of triton.ops.blocksparse.matmul('SDD').

from triton.

Could you please give some explanation on "how" to reduce the row in triton implementation?

from triton.

Thanks. Last question. Is this softmax specifically different from the softmax used for full attention, i.e. full spMM.

from triton.

Hmm, it should be the same. You can also pass dense masks since a block-triangular matrix is not triangular (the blocks on the diagonals are dense). There are example usages here https://github.com/ptillet/triton/blob/master/python/test/test_blocksparse.py#L47

from triton.