kempnerforge.model.masking

Attention mask builders.

One predicate covers every visibility rule self-attention implements:

visible(q, k) = same_document(q, k) AND (k <= q OR same_nonnegative_group(q, k)) AND valid_key(k)

build_attention_mask materializes it once per forward, in whichever form the configured backend consumes: a FlexAttention BlockMask under attention_backend="flex", a dense (batch, 1, seq_len, seq_len) bool tensor for F.scaled_dot_product_attention otherwise, or None when no term applies and plain causal attention is already exact. Attention and MoTAttention take that one object and derive no mask of their own, so a further visibility rule is a new argument here rather than a new parameter on every attention signature.

FlexAttention replaces the dense mask rather than computing it: the predicate is compiled into the attention kernel and fully-masked blocks are skipped rather than computed. For document packing the mask is block-diagonal and mostly zeros, so that is the difference between paying for the full S x S attention and paying only for the blocks inside a document. Key validity is deliberately not folded into the BlockMask; those batches take the dense form instead.

Functions

build_attention_mask(*, backend, batch, ...)

Build the one attention mask for a forward pass.

build_dense_attention_mask(*, batch, ...[, ...])

Dense boolean mask for F.scaled_dot_product_attention.

build_doc_causal_block_mask(doc_ids, device)

Block-diagonal causal mask: q attends to k iff same document and k <= q.

flex_attention_fn(compiled)

Return the flex_attention callable to use.

same_group_mask(group_ids)

Dense same_nonnegative_group term, for the dense mask path.

use_dynamic_flex_shapes(enabled)

Select the flex helpers compiled for a varying batch size, or the pinned-shape pair.
kempnerforge.model.masking.use_dynamic_flex_shapes(enabled)[source]

Select the flex helpers compiled for a varying batch size, or the pinned-shape pair.

Parameters:

enabled (bool)

Return type:

None

kempnerforge.model.masking.flex_attention_fn(compiled)[source]

Return the flex_attention callable to use.

Parameters:

compiled (bool) – Whether to return the torch.compile-wrapped kernel. True on CUDA; False on CPU, where the eager decomposition keeps unit tests off Inductor’s C++ codegen path. Note torch 2.11 has no CPU backward for FlexAttention, so the CPU path is forward-only.

Return type:

Callable[[…], Any]

kempnerforge.model.masking.build_attention_mask(*, backend, batch, seq_len, device, doc_ids=None, group_ids=None, key_padding_mask=None)[source]

Build the one attention mask for a forward pass.

The single entry point every caller uses, and the one place the ids are validated: it owns which backend form the semantics are expressed in, so callers pass ids and attention consumes one mask. backend="flex" returns a BlockMask whenever the predicate is expressible in it, which is every case except key validity; everything else returns the dense form.

The returned mask is square over the residual stream, so it is valid only for a forward whose keys are exactly its queries. Incremental decoding with a KVCache grows the key axis every step, which is why Transformer.forward refuses every mask input alongside kv_caches.

Parameters:

  • backend (str) – ModelConfig.attention_backend, "sdpa" or "flex".

  • batch (int) – Batch size of the residual stream.

  • seq_len (int) – Query and key length, the residual stream’s length.

  • device (torch.device) – Device to materialize the mask on. Any ids passed must already be on it.

  • doc_ids (torch.Tensor | None) – Optional per-token document ids, shape (batch, seq_len). Tokens attend only within their own document.

  • group_ids (torch.Tensor | None) – Optional per-token bidirectional-visibility group ids, shape (batch, seq_len); negative means no group. Tokens sharing a nonnegative id see each other regardless of order. All-negative ids give plain causal visibility but still build a mask; None keeps the unmasked path.

  • key_padding_mask (torch.Tensor | None) – Optional per-key validity mask, shape (batch, seq_len); True = attend, False = drop (e.g. the visual tokens of padded video frames).

Returns:

A BlockMask, a dense (batch, 1, seq_len, seq_len) bool tensor, or None when no term applies and attention is plain causal.

Raises:

ValueError – If any input’s shape is not exactly (batch, seq_len).

Return type:

torch.Tensor | BlockMask | None

kempnerforge.model.masking.same_group_mask(group_ids)[source]

Dense same_nonnegative_group term, for the dense mask path.

The counterpart of the group term inside build_doc_causal_block_mask.

Parameters:

group_ids (torch.Tensor) – Per-token group ids, shape (batch, seq_len); negative means no group.

Returns:

Bool tensor of shape (batch, 1, seq_len, seq_len), True where query and key share a group id and that id is nonnegative.

Return type:

torch.Tensor

kempnerforge.model.masking.build_dense_attention_mask(*, batch, seq_len, device, doc_ids=None, group_ids=None, key_padding_mask=None)[source]

Dense boolean mask for F.scaled_dot_product_attention.

True means the query may attend to the key. Starts causal, widens by the group term, then narrows by document isolation and key validity, so a group never reaches across a document and never revives an invalid key.

A query row left with no reachable key – the leading positions of an all-padded clip, say – would softmax over all -inf and yield NaN, so such rows are unmasked instead. Their outputs are discarded downstream, by output_slice or by the clip’s labels being all -100.

An explicit mask is not a FlashAttention-2 shape, so SDPA falls back to the mem-efficient or math kernel whenever one is returned. None keeps the is_causal fast path, which is why it is returned rather than an all-causal tensor.

Shapes are not checked here: build_attention_mask validates the ids before it hands them on.

Parameters:

  • batch (int) – Batch size of the residual stream.

  • seq_len (int) – Query and key length.

  • device (torch.device) – Device to materialize the causal term on. Any ids passed must already be on it.

  • doc_ids (torch.Tensor | None) – Optional per-token document ids, shape (batch, seq_len).

  • group_ids (torch.Tensor | None) – Optional per-token group ids, shape (batch, seq_len); negative means no group.

  • key_padding_mask (torch.Tensor | None) – Optional per-key validity mask, shape (batch, seq_len); True = attend.

Returns:

Bool tensor of shape (batch, 1, seq_len, seq_len), or None when every optional term is None and the mask would be plain causal.

Return type:

torch.Tensor | None

kempnerforge.model.masking.build_doc_causal_block_mask(doc_ids, device, group_ids=None)[source]

Block-diagonal causal mask: q attends to k iff same document and k <= q.

group_ids widens causality within caller-defined spans: visible(q, k) = same_document AND (k <= q OR same_nonnegative_group). Tokens with a negative group id stay purely causal, and document isolation always dominates, so a group id reused across documents never joins them. Without doc_ids the document term is dropped and each row is one document.

Built once per forward and shared by every layer. H=None broadcasts the mask over heads, which is what lets grouped-query attention pass enable_gqa=True instead of materializing repeated K/V heads, and what keeps the mask correct under tensor parallelism, where each rank holds only a shard of the heads.

seq_len must be at least FLEX_BLOCK_SIZE: below it an Inductor-compiled model silently leaks attention across document boundaries, so JobConfig.validate rejects that configuration outright. This kernel is not the culprit – it is exact at those lengths, as is the same model under backend="eager" or "aot_eager"; only Inductor codegen diverges.

Parameters:

  • doc_ids (torch.Tensor | None) – Optional per-token document ids, shape (batch, seq_len); None treats each row as one document.

  • device (torch.device) – Device on which to materialize the BlockMask.

  • group_ids (torch.Tensor | None) – Optional per-token bidirectional-visibility group ids, shape (batch, seq_len); negative means no group. None keeps the mask exactly causal-within-document. Ids are compared at full int64 width, unlike doc_ids, so any caller-defined id is usable.

Returns:

A BlockMask over (batch, seq_len, seq_len), head-broadcast.

Raises:

ValueError – If neither doc_ids nor group_ids is given.

Return type:

BlockMask