tensorplay.nn.attention.bias

Defines bias subclasses that work with scaled_dot_product_attention

A causal mask is the one attention bias common enough to deserve its own type. tensorplay.nn.attention.bias.CausalBias is that type: a non-materialized boolean mask that scaled_dot_product_attention (and its fused kernels, when eligible) interpret directly, without building the (L, S) mask tensor in memory.

Variants

tensorplay.nn.attention.bias.CausalVariant selects the alignment:

UPPER_LEFT — standard causal attention. Position i may attend to positions j <= i. The equivalent materialized mask is tensorplay.tril(ones(L, S, dtype=bool)).
LOWER_RIGHT — the allowed region is anchored to the lower-right corner instead, so the last query positions see the full history. With equal query and key lengths the two variants coincide; they differ when packing sequences of different lengths.

Constructing a bias

tensorplay.nn.attention.bias.causal_upper_left() and tensorplay.nn.attention.bias.causal_lower_right() take the query and key sequence lengths and return the matching bias:

import tensorplay as tp
import tensorplay.nn.functional as F
from tensorplay.nn.attention.bias import (
    causal_upper_left,
    CausalBias,
    CausalVariant,
)

q = tp.randn(1, 1, 8, 16)
k = tp.randn(1, 1, 8, 16)
v = tp.randn(1, 1, 8, 16)

bias = causal_upper_left(8, 8)
out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias)
print(out.shape)    # (1, 1, 8, 16)

The same object comes from the constructor form CausalBias(CausalVariant.UPPER_LEFT, seq_len_q, seq_len_kv); CausalVariant.LOWER_RIGHT gives the lower-right alignment. Because the bias is not a materialized tensor, it costs no L × S memory and is the mask form the flash-attention kernels prefer.

Flash-attention probes

The module re-exports the eligibility probes the dispatcher uses, so mask-related routing questions can be answered in one place:

`CausalBias`	A bias representing causal attention patterns.
`CausalVariant`	Enum for causal variants used in attention mechanisms.
`causal_upper_left`	Creates an upper-left triangular causal bias.
`causal_lower_right`	Creates a lower-right triangular causal bias.

`tensorplay.backends.cuda.is_flash_attention_available`	Check if TensorPlay was built with FlashAttention for scaled_dot_product_attention.
`tensorplay.backends.cuda.can_use_flash_attention`	Check if FlashAttention can be utilized in scaled_dot_product_attention.
`tensorplay.backends.cuda.can_use_efficient_attention`	Check if efficient_attention can be utilized in scaled_dot_product_attention.

tensorplay.nn.attention.bias

Variants

Constructing a bias

Flash-attention probes

Help improve this page

On this page