(self, hidden_size, num_heads, qkv_proj, o_proj, dev)
| 80 | """Multi-headed attention from 'Attention Is All You Need' paper""" |
| 81 | |
| 82 | def __init__(self, hidden_size, num_heads, qkv_proj, o_proj, dev): |
| 83 | super().__init__() |
| 84 | self.hidden_size = hidden_size |
| 85 | self.num_heads = num_heads |
| 86 | self.head_dim = hidden_size // num_heads |
| 87 | |
| 88 | if (self.head_dim * num_heads) != self.hidden_size: |
| 89 | raise ValueError( |
| 90 | f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" |
| 91 | f" and `num_heads`: {num_heads})." |
| 92 | ) |
| 93 | self.qkv_proj = qkv_proj |
| 94 | self.o_proj = o_proj |
| 95 | self.rotary_emb = QuantLlamaRotaryEmbedding( |
| 96 | self.head_dim, max_position_embeddings=2048, device=dev |
| 97 | ) |
| 98 | |
| 99 | def forward( |
| 100 | self, |
no test coverage detected