Forward pass of the attention module. Args: x (torch.Tensor): Input tensor. start_pos (int): Starting position for caching. freqs_cis (torch.Tensor): Precomputed frequency tensor. mask (torch.Tensor, optional): Attention mask tensor.
(
self,
x: torch.Tensor,
start_pos: int,
freqs_cis: torch.Tensor,
mask: Optional[torch.Tensor],
)
| 251 | ).cuda() |
| 252 | |
| 253 | def forward( |
| 254 | self, |
| 255 | x: torch.Tensor, |
| 256 | start_pos: int, |
| 257 | freqs_cis: torch.Tensor, |
| 258 | mask: Optional[torch.Tensor], |
| 259 | ): |
| 260 | """ |
| 261 | Forward pass of the attention module. |
| 262 | |
| 263 | Args: |
| 264 | x (torch.Tensor): Input tensor. |
| 265 | start_pos (int): Starting position for caching. |
| 266 | freqs_cis (torch.Tensor): Precomputed frequency tensor. |
| 267 | mask (torch.Tensor, optional): Attention mask tensor. |
| 268 | |
| 269 | Returns: |
| 270 | torch.Tensor: Output tensor after attention. |
| 271 | |
| 272 | """ |
| 273 | bsz, seqlen, _ = x.shape |
| 274 | xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) |
| 275 | |
| 276 | xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim) |
| 277 | xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim) |
| 278 | xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim) |
| 279 | |
| 280 | xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis) |
| 281 | |
| 282 | self.cache_k = self.cache_k.to(xq) |
| 283 | self.cache_v = self.cache_v.to(xq) |
| 284 | |
| 285 | self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk |
| 286 | self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv |
| 287 | |
| 288 | keys = self.cache_k[:bsz, : start_pos + seqlen] |
| 289 | values = self.cache_v[:bsz, : start_pos + seqlen] |
| 290 | |
| 291 | # repeat k/v heads if n_kv_heads < n_heads |
| 292 | keys = repeat_kv( |
| 293 | keys, self.n_rep |
| 294 | ) # (bs, cache_len + seqlen, n_local_heads, head_dim) |
| 295 | values = repeat_kv( |
| 296 | values, self.n_rep |
| 297 | ) # (bs, cache_len + seqlen, n_local_heads, head_dim) |
| 298 | |
| 299 | xq = xq.transpose(1, 2) # (bs, n_local_heads, seqlen, head_dim) |
| 300 | keys = keys.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim) |
| 301 | values = values.transpose( |
| 302 | 1, 2 |
| 303 | ) # (bs, n_local_heads, cache_len + seqlen, head_dim) |
| 304 | scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim) |
| 305 | if mask is not None: |
| 306 | scores = scores + mask # (bs, n_local_heads, seqlen, cache_len + seqlen) |
| 307 | scores = F.softmax(scores.float(), dim=-1).type_as(xq) |
| 308 | output = torch.matmul(scores, values) # (bs, n_local_heads, seqlen, head_dim) |
| 309 | output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1) |
| 310 | return self.wo(output) |
nothing calls this directly
no test coverage detected