# tensorplay.nn.attention.create_block_mask

Source: https://www.tensorplay.cn/docs/generated/tensorplay.nn.attention.create_block_mask.html

```python
tensorplay.nn.attention.create_block_mask(mask_mod: _mask_mod_signature, B: int | None, H: int | None, Q_LEN: int, KV_LEN: int, device: DeviceLikeType | None = None, BLOCK_SIZE: int | tuple[int, int] = 128, _compile=False, separate_full_blocks: bool = True, compute_dq_write_order: bool = False, dq_kv_order: bool = True) → BlockMask
```

This function creates a block mask tuple from a mask_mod function.

Parameters:

- mask_mod (Callable) – mask_mod function. This is a callable that defines the masking pattern for the attention mechanism. It takes four arguments: b (batch size), h (number of heads), q_idx (query index), and kv_idx (key/value index). It should return a boolean tensor indicating which attention connections are allowed (True) or masked out (False).

- B ([int](https://docs.python.org/3/builtins/functions.html#int)) – Batch size.

- H ([int](https://docs.python.org/3/builtins/functions.html#int)) – Number of query heads.

- Q_LEN ([int](https://docs.python.org/3/builtins/functions.html#int)) – Sequence length of query.

- KV_LEN ([int](https://docs.python.org/3/builtins/functions.html#int)) – Sequence length of key/value.

- device ([str](https://docs.python.org/3/builtins/stdtypes.html#str)) – Device to run the mask creation on.

- BLOCK_SIZE ([int](https://docs.python.org/3/builtins/functions.html#int) or [tuple](https://docs.python.org/3/builtins/stdtypes.html#tuple)[[int](https://docs.python.org/3/builtins/functions.html#int), [int](https://docs.python.org/3/builtins/functions.html#int)]) – Block size for the block mask. If a single int is provided it is used for both query and key/value.

- separate_full_blocks ([bool](https://docs.python.org/3/builtins/functions.html#bool)) – If True, fully unmasked blocks are stored separately so kernels can skip mask_mod on those blocks. If False, all non-empty blocks are stored as partial blocks and mask_mod is applied to every block.

- compute_dq_write_order ([bool](https://docs.python.org/3/builtins/functions.html#bool)) – If True, precompute dQ write-order metadata needed by deterministic block-sparse backward.

- dq_kv_order ([bool](https://docs.python.org/3/builtins/functions.html#bool)) – KV-column scheduler order used for deterministic dQ accumulation when compute_dq_write_order is True. False means ascending n-block order and True means descending/SPT order. Explicit tensor schedules are not supported by create_block_mask yet; they are supported by BlockMask.from_kv_blocks for callers that provide precomputed write-order metadata directly.

Returns:

A BlockMask object that contains the block mask information.

Return type:

[BlockMask](/docs/generated/tensorplay.nn.attention.BlockMask.html#tensorplay.nn.attention.BlockMask)

Example Usage:

```
def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

block_mask = create_block_mask(causal_mask, 1, 1, 8192, 8192, device="cuda")
query = tensorplay.randn(1, 1, 8192, 64, device="cuda", dtype=tensorplay.float16)
key = tensorplay.randn(1, 1, 8192, 64, device="cuda", dtype=tensorplay.float16)
value = tensorplay.randn(1, 1, 8192, 64, device="cuda", dtype=tensorplay.float16)
output = omni_attention(query, key, value, block_mask=block_mask)
```