llama_xformers_attention.py

import torch
import torch.nn as nn

from typing import Optional, Tuple

from transformers.models.llama.modeling_llama import LlamaAttention, apply_rotary_pos_emb, repeat_kv

import xformers
from xformers.ops.fmha import (
 memory_efficient_attention,
)

class LlamaXFormersAttention(LlamaAttention):
 def forward(
 self,
 hidden_states: torch.Tensor,
 attention_mask: Optional[torch.Tensor] = None,
 position_ids: Optional[torch.LongTensor] = None,
 past_key_value: Optional[Tuple[torch.Tensor]] = None,
 output_attentions: bool = False,
 use_cache: bool = False,
 **kwargs,
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
 if "padding_mask" in kwargs:
 warnings.warn(
 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
 )

 bsz, q_len, _ = hidden_states.size()

 query_states = self.q_proj(hidden_states)
 key_states = self.k_proj(hidden_states)
 value_states = self.v_proj(hidden_states)

 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

 kv_seq_len = key_states.shape[-2]
 if past_key_value is not None:
 kv_seq_len += past_key_value[0].shape[-2]
 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

 if past_key_value is not None:
 # reuse k, v, self_attention
 key_states = torch.cat([past_key_value[0], key_states], dim=2)
 value_states = torch.cat([past_key_value[1], value_states], dim=2)

 past_key_value = (key_states, value_states) if use_cache else None

 key_states = repeat_kv(key_states, self.num_key_value_groups)
 value_states = repeat_kv(value_states, self.num_key_value_groups)

 dtype = query_states.dtype

 query_states = query_states.transpose(1, 2)
 key_states = key_states.transpose(1, 2)
 value_states = value_states.transpose(1, 2)

 # copied from https://github.com/oobabooga/text-generation-webui/pull/950/files
 # This is a nasty hack. We know attention_mask in transformers is either LowerTriangular or all Zeros.
 # We therefore check if one element in the upper triangular portion is zero. If it is, then the mask is all zeros.
 if attention_mask is None or attention_mask[0, 0, 0, 1] == 0:
 # input and output should be of form (bsz, q_len, num_heads, head_dim)
 attn_output = memory_efficient_attention(query_states, key_states, value_states, attn_bias=None)
 else:
 # input and output should be of form (bsz, q_len, num_heads, head_dim)
 attn_output = memory_efficient_attention(query_states, key_states, value_states, attn_bias=xformers.ops.LowerTriangularMask())
 attn_weights = None

 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

 attn_output = self.o_proj(attn_output)

 if not output_attentions:
 attn_weights = None

 return attn_output, attn_weights, past_key_value