class_docstr: Multi-head attention module.
class_docstr_tok: ["Multi-head","attention","module","."]
docstr: 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.
  
  Returns:
      torch.Tensor: Output tensor after attention.
docstr_tok: ["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",".","Returns",":","torch.Tensor",":","Output","tensor","after","attention","."]
code: def forward(
      self,
      x: torch.Tensor,
      start_pos: int,
      freqs_cis: torch.Tensor,
      mask: Optional[torch.Tensor],
  ):
      """
      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.
  
      Returns:
          torch.Tensor: Output tensor after attention.
  
      """
      bsz, seqlen, _ = x.shape
      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
  
      xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
      xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
      xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
  
      xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
  
      self.cache_k = self.cache_k.to(xq)
      self.cache_v = self.cache_v.to(xq)
  
      self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
      self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
  
      keys = self.cache_k[:bsz, : start_pos + seqlen]
      values = self.cache_v[:bsz, : start_pos + seqlen]
  
      # repeat k/v heads if n_kv_heads < n_heads
      keys = repeat_kv(
          keys, self.n_rep
      )  # (bs, cache_len + seqlen, n_local_heads, head_dim)
      values = repeat_kv(
          values, self.n_rep
      )  # (bs, cache_len + seqlen, n_local_heads, head_dim)
  
      xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
      keys = keys.transpose(
          1, 2
      )  # (bs, n_local_heads, cache_len + seqlen, head_dim)
      values = values.transpose(
          1, 2
      )  # (bs, n_local_heads, cache_len + seqlen, head_dim)
      scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(
          self.head_dim
      )
      if mask is not None:
          scores = (
              scores + mask
          )  # (bs, n_local_heads, seqlen, cache_len + seqlen)
      scores = F.softmax(scores.float(), dim=-1).type_as(xq)
      output = torch.matmul(
          scores, values
      )  # (bs, n_local_heads, seqlen, head_dim)
      output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
      return self.wo(output)
  
imports: import math
  from dataclasses import dataclass
  from typing import Optional, Tuple
  import fairscale.nn.model_parallel.initialize
  import torch
  import torch.nn.functional
  from fairscale.nn.model_parallel.layers import ColumnParallelLinear, ParallelEmbedding, RowParallelLinear
  from torch import nn
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(x, start_pos, freqs_cis, mask) and returns: self