BlockMask

class tensorplay.nn.attention.BlockMask(seq_lengths: tuple[int, int], kv_num_blocks: ~tensorplay.Tensor, kv_indices: ~tensorplay.Tensor, full_kv_num_blocks: ~tensorplay.Tensor | None, full_kv_indices: ~tensorplay.Tensor | None, q_num_blocks: ~tensorplay.Tensor | None, q_indices: ~tensorplay.Tensor | None, full_q_num_blocks: ~tensorplay.Tensor | None, full_q_indices: ~tensorplay.Tensor | None, BLOCK_SIZE: tuple[int, int] = (128, 128), mask_mod: ~collections.abc.Callable[[~tensorplay.Tensor, ~tensorplay.Tensor, ~tensorplay.Tensor, ~tensorplay.Tensor], ~tensorplay.Tensor] = <function noop_mask>, *, dq_write_order: ~tensorplay.Tensor | None = None, dq_write_order_full: ~tensorplay.Tensor | None = None, dq_kv_order: ~tensorplay.Tensor | None = None, dq_kv_order_spt: bool | None = None)[source]

BlockMask is our format for representing a block-sparse attention mask. It is somewhat of a cross in-between BCSR and a non-sparse format.

Basics

A block-sparse mask means that instead of representing the sparsity of individual elements in the mask, a KV_BLOCK_SIZE x Q_BLOCK_SIZE block is considered sparse only if every element within that block is sparse. This aligns well with hardware, which generally expects to perform contiguous loads and computation.

This format is primarily optimized for 1. simplicity, and 2. kernel efficiency. Notably, it is not optimized for size, as this mask is always reduced by a factor of KV_BLOCK_SIZE * Q_BLOCK_SIZE. If the size is a concern, the tensors can be reduced in size by increasing the block size.

The essentials of our format are:

num_blocks_in_row: Tensor[ROWS]: Describes the number of blocks present in each row.

col_indices: Tensor[ROWS, MAX_BLOCKS_IN_COL]: col_indices[i] is the sequence of block positions for row i. The values of this row after col_indices[i][num_blocks_in_row[i]] are undefined.

For example, to reconstruct the original tensor from this format:

dense_mask = tensorplay.zeros(ROWS, COLS)
for row in range(ROWS):
    for block_idx in range(num_blocks_in_row[row]):
        dense_mask[row, col_indices[row, block_idx]] = 1

Notably, this format makes it easier to implement a reduction along the rows of the mask.

Details

The basics of our format require only kv_num_blocks and kv_indices. The primary block-sparse layout is represented by up to 4 tensor pairs:

1. (kv_num_blocks, kv_indices): Used for the forwards pass of attention, as we reduce along the KV dimension.

2. [OPTIONAL] (full_kv_num_blocks, full_kv_indices): This is optional and purely an optimization. As it turns out, applying masking to every block is quite expensive! If we specifically know which blocks are “full” and don’t require masking at all, then we can skip applying mask_mod to these blocks. This requires the user to split out a separate mask_mod from the score_mod. For causal masks, this is about a 15% speedup.

3. [GENERATED] (q_num_blocks, q_indices): Required for the backwards pass, as computing dKV requires iterating along the mask along the Q dimension. These are autogenerated from 1.

4. [GENERATED] (full_q_num_blocks, full_q_indices): Same as above, but for the backwards pass. These are autogenerated from 2.

Additional optional tensors may carry deterministic dQ metadata for block-sparse backward:

5. [OPTIONAL] dq_write_order: Write-order metadata for partial blocks. This is produced by create_block_mask when compute_dq_write_order=True, or passed directly to BlockMask.from_kv_blocks by callers that precompute it.

6. [OPTIONAL] dq_write_order_full: Write-order metadata for full blocks, produced or passed the same way as dq_write_order.

7. [OPTIONAL] dq_kv_order: Explicit KV scheduler order used to produce the write-order metadata. create_block_mask currently accepts a boolean dq_kv_order; BlockMask.from_kv_blocks also accepts a tensor for callers that provide precomputed write-order metadata directly.

as_tuple(flatten: Literal[True] = True) → tuple[int, int, Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, bool | None, int, int, Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]][source]

as_tuple(flatten: Literal[False]) → tuple[tuple[int, int], Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, Tensor | None, bool | None, tuple[int, int], Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]]

Returns a tuple of the attributes of the BlockMask.

Parameters:

flatten (bool) – If True, it will flatten the tuple of (KV_BLOCK_SIZE, Q_BLOCK_SIZE)

classmethod from_kv_blocks(kv_num_blocks: Tensor, kv_indices: Tensor, full_kv_num_blocks: Tensor | None = None, full_kv_indices: Tensor | None = None, BLOCK_SIZE: int | tuple[int, int] = 128, mask_mod: Callable[[Tensor, Tensor, Tensor, Tensor], Tensor] | None = None, seq_lengths: tuple[int, int] | None = None, compute_q_blocks: bool = True, *, dq_write_order: Tensor | None = None, dq_write_order_full: Tensor | None = None, dq_kv_order: Tensor | bool | None = None) → Self[source]

Creates a BlockMask instance from key-value block information.

Parameters:

• kv_num_blocks (Tensor) – Number of kv_blocks in each Q_BLOCK_SIZE row tile.

• kv_indices (Tensor) – Indices of key-value blocks in each Q_BLOCK_SIZE row tile.

• full_kv_num_blocks (Optional[Tensor]) – Number of full kv_blocks in each Q_BLOCK_SIZE row tile.

• full_kv_indices (Optional[Tensor]) – Indices of full key-value blocks in each Q_BLOCK_SIZE row tile.

• BLOCK_SIZE (Union[int, tuple[int, int]]) – Size of KV_BLOCK_SIZE x Q_BLOCK_SIZE tiles.

• mask_mod (Optional[Callable]) – Function to modify the mask.

• dq_write_order (Optional[Tensor]) – Precomputed deterministic dQ write-order metadata.

• dq_write_order_full (Optional[Tensor]) – Precomputed deterministic dQ write-order metadata for full blocks.

• dq_kv_order (Optional[Union[Tensor, bool]]) – KV-column scheduler order used to produce dq_write_order. A bool selects a built-in order; a tensor gives an explicit scheduler-rank to n-block permutation.

Returns:

Instance with full Q information generated via _transposed_ordered

Return type:

BlockMask

Raises:

• RuntimeError – If kv_indices has < 2 dimensions.

• AssertionError – If only one of full_kv_* args is provided.

numel() → int[source]

Returns the number of elements (not accounting for sparsity) in the mask.

sparsity() → float[source]

Computes the percentage of blocks that are sparse (i.e. not computed)

to(device: device | str) → BlockMask[source]

Moves the BlockMask to the specified device.

Parameters:

device (tensorplay.device or str) – The target device to move the BlockMask to. Can be a device object or a string (e.g., ‘cpu’, ‘cuda:0’).

Returns:

A new BlockMask instance with all tensor components moved to the specified device.

Return type:

BlockMask

Note

This method does not modify the original BlockMask in-place. Instead, it returns a new BlockMask instance where individual tensor attributes may or may not be moved to the specified device, depending on their current device placement.

to_dense() → Tensor[source]

Returns a dense block that is equivalent to the block mask.

to_string(grid_size: int | tuple[int, int] = (20, 20), limit: int = 4) → str[source]

Returns a string representation of the block mask. Quite nifty.

If grid_size is -1, prints out an uncompressed version. Warning, it can be quite big!