Improvements in attention_forward functions

Author: mseeger

src/transformers/integrations/flash_attention.py

@@ -20,9 +20,16 @@ def flash_attention_forward(
     sliding_window: Optional[int] = None,
     softcap: Optional[float] = None,
     **kwargs,
-) -> Tuple[torch.Tensor, None]:
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    `query` has shape `(batch, num_heads, q_len, head_dim)`, while `key`,
+    `value` have shape `(batch, num_key_value_groups, kv_len, head_dim)`. Here,
+    `num_key_value_groups <= num_heads` and
+    `num_heads % num_key_value_groups == 0`.
+
+    """
     # This is before the transpose
-    seq_len = query.shape[2]
+    q_len = query.shape[2]
 
     # FA2 uses non-transposed inputs
     query = query.transpose(1, 2)
@@ -52,7 +59,7 @@ def flash_attention_forward(
         key,
         value,
         attention_mask,
-        query_length=seq_len,
+        query_length=q_len,
         is_causal=module.is_causal,
         dropout=dropout,
         softmax_scale=scaling,
@@ -62,5 +69,6 @@ def flash_attention_forward(
         target_dtype=target_dtype,
         **kwargs,
     )
+    # attn_output: (batch, q_len, num_heads, head_dim)
 
     return attn_output, None

src/transformers/integrations/flex_attention.py

@@ -20,6 +20,13 @@ def flex_attention_forward(
     head_mask: Optional[torch.Tensor] = None,
     **kwargs,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    `query` has shape `(batch, num_heads, q_len, head_dim)`, while `key`,
+    `value` have shape `(batch, num_key_value_groups, kv_len, head_dim)`. Here,
+    `num_key_value_groups <= num_heads` and
+    `num_heads % num_key_value_groups == 0`.
+
+    """
     causal_mask = attention_mask
     if causal_mask is not None:
         causal_mask = causal_mask[:, :, :, : key.shape[-2]]

src/transformers/integrations/sdpa_attention.py

@@ -25,40 +25,60 @@ def sdpa_attention_forward(
     scaling: Optional[float] = None,
     is_causal: Optional[bool] = None,
     **kwargs,
-) -> Tuple[torch.Tensor, None]:
-    if hasattr(module, "num_key_value_groups"):
-        key = repeat_kv(key, module.num_key_value_groups)
-        value = repeat_kv(value, module.num_key_value_groups)
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    `query` has shape `(batch, num_heads, q_len, head_dim)`, while `key`,
+    `value` have shape `(batch, num_key_value_groups, kv_len, head_dim)`. Here,
+    `num_key_value_groups <= num_heads` and
+    `num_heads % num_key_value_groups == 0`.
+
+    https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html#torch-nn-functional-scaled-dot-product-attention
+    `scaled_dot_product_attention` is supposed to support
+    `num_key_value_groups < num_heads`. But at least up to PyTorch 2.5.1, this
+    does not seem to be supported.
 
+    """
+    assert query.ndim == key.ndim == value.ndim == 4
+    _, num_heads, q_len, _ = query.shape
+    _, num_key_value_groups, kv_len, _ = key.shape
+    assert query.shape[0] == key.shape[0] == value.shape[0]  # batch_size
+    assert value.shape[1] == num_key_value_groups and value.shape[2] == kv_len
+    assert num_heads % num_key_value_groups == 0 and num_heads >= num_key_value_groups
+    q_per_kv = num_heads // num_key_value_groups
     causal_mask = attention_mask
     if attention_mask is not None:
-        causal_mask = causal_mask[:, :, :, : key.shape[-2]]
-
-    # SDPA with memory-efficient backend is bugged with non-contiguous inputs and custom attn_mask for some torch versions
-    # Reference: https://github.com/pytorch/pytorch/issues/112577.
-    query = query.contiguous()
-    key = key.contiguous()
-    value = value.contiguous()
-
+        causal_mask = causal_mask[:, :, :, :kv_len]
     # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
     # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
     if is_causal is None:
-        is_causal = causal_mask is None and query.shape[2] > 1
+        is_causal = causal_mask is None and q_len > 1
 
     # Shapes (e.g. query.shape[2]) are tensors during jit tracing, resulting in `is_causal` being a tensor.
     # We convert it to a bool for the SDPA kernel that only accepts bools.
     if torch.jit.is_tracing() and isinstance(is_causal, torch.Tensor):
         is_causal = is_causal.item()
 
+    # SDPA with memory-efficient backend is bugged with non-contiguous inputs and custom attn_mask for some torch versions
+    # Reference: https://github.com/pytorch/pytorch/issues/112577.
+    query = query.contiguous()
+    if q_per_kv == 1:
+        key = key.contiguous()
+        value = value.contiguous()
+    else:
+        # TODO: Once PyTorch SDPA supports q_per_kv > 1, this should be
+        # removed!
+        key = repeat_kv(key, n_rep=q_per_kv).contiguous()
+        value = repeat_kv(value, n_rep=q_per_kv).contiguous()
+
     attn_output = torch.nn.functional.scaled_dot_product_attention(
-        query,
-        key,
-        value,
+        query=query,
+        key=key,
+        value=value,
         attn_mask=causal_mask,
         dropout_p=dropout,
         scale=scaling,
         is_causal=is_causal,
     )
     attn_output = attn_output.transpose(1, 2).contiguous()
-
+    # attn_output: (batch, q_len, num_heads, head_dim)
     return attn_output, None

src/transformers/models/aria/configuration_aria.py

@@ -22,9 +22,14 @@
 
 from ...configuration_utils import PretrainedConfig
 from ...modeling_rope_utils import rope_config_validation
+from ...utils.import_utils import is_torch_available
 from ..auto import CONFIG_MAPPING, AutoConfig
 
 
+if is_torch_available():
+    pass
+
+
 class AriaTextConfig(PretrainedConfig):
     r"""
     This class handles the configuration for the text component of the Aria model.

src/transformers/models/aria/image_processing_aria.py

@@ -37,6 +37,11 @@
     validate_preprocess_arguments,
 )
 from ...utils import TensorType
+from ...utils.import_utils import is_torch_available
+
+
+if is_torch_available():
+    pass
 
 
 def divide_to_patches(image: np.array, patch_size: int, input_data_format) -> List[np.array]:

src/transformers/models/aria/modeling_aria.py

@@ -45,6 +45,7 @@
 
 if is_torch_available():
     import torch
+    import torch.nn.functional as F
     from torch import nn
 
 
@@ -346,7 +347,7 @@ def forward(self, permuted_tokens, tokens_per_expert):
         """
         fc1_output = self.fc1(permuted_tokens, tokens_per_expert)
         projection, gate = torch.chunk(fc1_output, 2, dim=-1)
-        fc1_output = nn.functional.silu(projection) * gate
+        fc1_output = F.silu(projection) * gate
         fc2_output = self.fc2(fc1_output, tokens_per_expert)
         return fc2_output
 
@@ -395,7 +396,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # Top K Routing
         logits = self.router(hidden_states)
         top_logits, top_indices = torch.topk(logits, k=self.config.moe_topk, dim=1)
-        scores = nn.functional.softmax(top_logits, dim=-1)
+        scores = F.softmax(top_logits, dim=-1)
 
         original_dtype = top_indices.dtype
 
@@ -465,16 +466,61 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     return q_embed, k_embed
 
 
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """
-    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
-    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
-    """
-    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-    if n_rep == 1:
-        return hidden_states
-    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
-    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+def _attention_compute_scores(
+    query: torch.Tensor,
+    key: torch.Tensor,
+) -> torch.Tensor:
+    nh_q = query.shape[1]
+    nh_k = key.shape[1]
+    # - query: (bs, nh_q, T_q, hs)
+    # - key: (bs, nh_k, T_k, hs)
+    q_per_kv = nh_q // nh_k
+    key_transposed = key.mT  # (bs, nh_k, hs, T_k)
+    if q_per_kv == 1:
+        return query @ key_transposed
+    else:
+        assert q_per_kv > 1
+        if nh_k > 1:
+            q_shape = query.shape[:1] + (nh_k, q_per_kv) + query.shape[2:]
+            _query = query.view(*q_shape)
+            key_transposed = key_transposed.unsqueeze(2)
+        else:
+            _query = query
+        # At this point:
+        # - _query: (bs, nh_k, q_per_kv, T_q, hs)
+        # - key_transposed: (bs, nh_k, 1, hs, T_k)
+        # - scores: (bs, nh_k, q_per_kv, T_q, T_k) -> (bs, nh_q, T_q, T_k)
+        scores = torch.matmul(_query, key_transposed)
+        s_shape = query.shape[:-1] + (key.shape[2],)
+        return scores.view(*s_shape)
+
+
+def _attention_compute_weighted_values(
+    scores: torch.Tensor,
+    value: torch.Tensor,
+) -> torch.Tensor:
+    nh_q = scores.shape[1]
+    nh_k = value.shape[1]
+    # - scores: (bs, nh_q, T_q, T_k)
+    # - value: (bs, nh_k, T_k, hs)
+    q_per_kv = nh_q // nh_k
+    if q_per_kv == 1:
+        return scores @ value
+    else:
+        if nh_k > 1:
+            s_shape = scores.shape[:1] + (nh_k, q_per_kv) + scores.shape[2:]
+            _scores = scores.view(*s_shape)
+            _value = value.unsqueeze(2)
+        else:
+            _scores = scores
+            _value = value
+        # At this point:
+        # - _scores: (bs, nh_k, q_per_kv, T_q, T_k)
+        # - _value: (bs, nh_k, 1, T_k, hs)
+        # - result: (bs, nh_k, q_per_kv, T_q, hs) -> (bs, nh_q, T_q, hs)
+        result = torch.matmul(_scores, _value)
+        r_shape = scores.shape[:-1] + (value.shape[-1],)
+        return result.view(*r_shape)
 
 
 def eager_attention_forward(
@@ -486,19 +532,32 @@ def eager_attention_forward(
     scaling: float,
     dropout: float = 0.0,
     **kwargs,
-):
-    key_states = repeat_kv(key, module.num_key_value_groups)
-    value_states = repeat_kv(value, module.num_key_value_groups)
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """
+    `query` has shape `(batch, num_heads, q_len, head_dim)`, while `key`,
+    `value` have shape `(batch, num_key_value_groups, kv_len, head_dim)`. Here,
+    `num_key_value_groups <= num_heads` and
+    `num_heads % num_key_value_groups == 0`.
 
-    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    """
+    assert query.ndim == key.ndim == value.ndim == 4
+    _, num_heads, q_len, _ = query.shape
+    _, num_key_value_groups, kv_len, _ = key.shape
+    assert query.shape[0] == key.shape[0] == value.shape[0]  # batch_size
+    assert value.shape[1] == num_key_value_groups and value.shape[2] == kv_len
+    assert num_heads % num_key_value_groups == 0 and num_heads >= num_key_value_groups
+
+    attn_weights = _attention_compute_scores(query, key) * scaling
     if attention_mask is not None:
-        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        causal_mask = attention_mask[:, :, :, :kv_len]
         attn_weights = attn_weights + causal_mask
 
-    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
-    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
-    attn_output = torch.matmul(attn_weights, value_states)
+    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = F.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = _attention_compute_weighted_values(attn_weights, value)
     attn_output = attn_output.transpose(1, 2).contiguous()
+    # attn_output: (batch, q_len, num_heads, head_dim)
+    # attn_weights: (batch, num_heads, q_len, kv_len)
 
     return attn_output, attn_weights
 
@@ -969,6 +1028,7 @@ def forward(
                     use_cache,
                     cache_position,
                     position_embeddings,
+                    **flash_attn_kwargs,
                 )
             else:
                 layer_outputs = decoder_layer(

src/transformers/models/aria/modular_aria.py

@@ -66,6 +66,7 @@
 
 if is_torch_available():
     import torch
+    import torch.nn.functional as F
     from torch import nn
 
 
@@ -1094,7 +1095,7 @@ def forward(self, permuted_tokens, tokens_per_expert):
         """
         fc1_output = self.fc1(permuted_tokens, tokens_per_expert)
         projection, gate = torch.chunk(fc1_output, 2, dim=-1)
-        fc1_output = nn.functional.silu(projection) * gate
+        fc1_output = F.silu(projection) * gate
         fc2_output = self.fc2(fc1_output, tokens_per_expert)
         return fc2_output
 
@@ -1143,7 +1144,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         # Top K Routing
         logits = self.router(hidden_states)
         top_logits, top_indices = torch.topk(logits, k=self.config.moe_topk, dim=1)
-        scores = nn.functional.softmax(top_logits, dim=-1)
+        scores = F.softmax(top_logits, dim=-1)
 
         original_dtype = top_indices.dtype
 

src/transformers/models/aria/processing_aria.py

@@ -25,9 +25,14 @@
 from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
 from ...tokenization_utils import PreTokenizedInput, TextInput
 from ...utils import TensorType
+from ...utils.import_utils import is_torch_available
 from ..auto import AutoTokenizer
 
 
+if is_torch_available():
+    pass
+
+
 class AriaProcessorKwargs(ProcessingKwargs, total=False):
     _defaults = {
         "text_kwargs": {