class_docstr: Implement the scaled dot product attention with softmax.
  Arguments
  ---------
      softmax_scale: The temperature to use for the softmax attention.
                    (default: 1/sqrt(d_keys) where d_keys is computed at
                    runtime)
      attention_dropout: The dropout rate to apply to the attention
                         (default: 0.0)
docstr: Implements the multihead softmax attention.
  Arguments
  ---------
      qkv: The tensor containing the query, key, and value. (B, S, 3, H, D)
docstr_tok: ["Implements","the","multihead","softmax","attention",".","Arguments","--","--","--","--","-","qkv",":","The","tensor","containing","the","query",",","key",",","and","value",".","(","B",",","S",",","3",",","H",",","D",")"]
code: def forward(self, qkv):
      """Implements the multihead softmax attention.
      Arguments
      ---------
          qkv: The tensor containing the query, key, and value. (B, S, 3, H, D)
      """
      assert qkv.dtype in [torch.float16, torch.bfloat16]
      assert qkv.is_cuda
      batch_size, seqlen = qkv.shape[0], qkv.shape[1]
      qkv = rearrange(qkv, "b s ... -> (b s) ...")
      max_s = seqlen
      cu_seqlens = torch.arange(
          0,
          (batch_size + 1) * seqlen,
          step=seqlen,
          dtype=torch.int32,
          device=qkv.device,
      )
      output = flash_attn_unpadded_qkvpacked_func(
          qkv,
          cu_seqlens,
          max_s,
          self.dropout_p if self.training else 0.0,
          softmax_scale=self.softmax_scale,
          causal=self.causal,
      )
      output = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
      return output
  
imports: import math
  from functools import partial
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
label_desc: Use image node_id 2 for calling the FlashSelfAttention obj's underlying member method code with example usage: obj.forward(qkv) and returns: output