Skip to content

joint_attention_base

joint_attention_base

Base class for dual-stream attention mechanisms.

Provides shared K/V concatenation, GQA expansion, dual SDPA, mask building, and tensor reshaping used by both full and precomputed dual-stream attention variants.

References

Esser et al. "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis" https://arxiv.org/abs/2403.03206

JointAttentionBase 

JointAttentionBase(number_of_heads, number_of_key_value_heads, head_dimension, dropout=0.0)

Bases: Module

Shared dual-stream attention computation.

Subclasses handle projection (full vs precomputed primary) and call _joint_sdpa with the projected Q/K/V tensors.

Shape notation

B: batch size S: primary sequence length T: secondary sequence length H: number of query heads KV_H: number of key/value heads D_head: per-head dimension

Source code in src/versatil/models/layers/transformer/attention/joint_attention_base.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def __init__(
    self,
    number_of_heads: int,
    number_of_key_value_heads: int,
    head_dimension: int,
    dropout: float = 0.0,
):
    super().__init__()
    if number_of_heads <= 0:
        raise ValueError(
            f"number_of_heads must be positive, got {number_of_heads}."
        )
    if number_of_key_value_heads <= 0:
        raise ValueError(
            "number_of_key_value_heads must be positive, "
            f"got {number_of_key_value_heads}."
        )
    if number_of_heads % number_of_key_value_heads != 0:
        raise ValueError(
            f"number_of_heads ({number_of_heads}) must be divisible by "
            f"number_of_key_value_heads ({number_of_key_value_heads})."
        )
    self.number_of_heads = number_of_heads
    self.number_of_key_value_heads = number_of_key_value_heads
    self.head_dimension = head_dimension
    self.group_size = number_of_heads // number_of_key_value_heads
    self.dropout = dropout