| import torch | |
| from torch.nn.attention.flex_attention import create_block_mask | |
| def build_doc_block_mask(segment_ids: torch.Tensor, seq_len: int, device): | |
| """Build the FlexAttention block mask for causal + same-document attention. | |
| segment_ids: (B, T) int32 tensor from the dataset. | |
| A query at position q attends to a key at position kv only when | |
| kv <= q (causal) AND both positions belong to the same packed document. | |
| """ | |
| B = segment_ids.shape[0] | |
| def doc_mask_mod(b, h, q_idx, kv_idx): | |
| causal = q_idx >= kv_idx | |
| same_doc = segment_ids[b, q_idx] == segment_ids[b, kv_idx] | |
| return causal & same_doc | |
| # H=None means one mask is shared across all heads | |
| return create_block_mask( | |
| doc_mask_mod, B=B, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=device | |
| ) |