class_docstr: Base mixin for all attention modules. The interface represented by this module
  is broadly:
  
  output = attention(query_key_value,
                     Optional[kv_cache],
                     Optional[attention_mask],
                     Optional[attention_bias])
class_docstr_tok: ["Base","mixin","for","all","attention","modules",".","The","interface","represented","by","this","module","is","broadly",":","output","=","attention","(","query_key_value",",","Optional","[","kv_cache","]",",","Optional","[","attention_mask","]",",","Optional","[","attention_bias","]",")"]
docstr: Parameters:
      q_k_v (torch.Tensor): Query, key, and value tensors. Expected shape is:
          [
              batch,
              seq_len,
              2 * self._config.n_heads_kv + self._config.n_heads_q,
              self._config.head_size
          ].
      kv_cache (Optional[torch.Tensor]): Key and value cache tensor. Expected shape is
          [
              2,
              batch,
              kv_cache_len,
              self._config.n_heads_kv,
              self._config.head_size
          ]. If None, cache is disabled. The `kv_cache_len` dimension does not need to
          be contiguous (it should expand stride by `max_out_tokens`).
      batch (RaggedBatchWrapper): Ragged batch metadata.
      attention_mask (Optional[torch.Tensor]): Attention mask tensor. If None, masking is
          disabled. This will defer to the config in the case of conflicting information.
          This means if the config class is implying causal attention, the mask will be ignored.
      attention_bias (Optional[torch.Tensor]): Attention bias tensor. If None, bias is disabled.
docstr_tok: ["Parameters",":","q_k_v","(","torch.Tensor",")",":","Query",",","key",",","and","value","tensors",".","Expected","shape","is",":","[","batch",",","seq_len",",","2","*","self._config.n_heads_kv","+","self._config.n_heads_q",",","self._config.head_size","]",".","kv_cache","(","Optional","[","torch.Tensor","]",")",":","Key","and","value","cache","tensor",".","Expected","shape","is","[","2",",","batch",",","kv_cache_len",",","self._config.n_heads_kv",",","self._config.head_size","]",".","If","None",",","cache","is","disabled",".","The","`","kv_cache_len","`","dimension","does","not","need","to","be","contiguous","(","it","should","expand","stride","by","`","max_out_tokens","`",")",".","batch","(","RaggedBatchWrapper",")",":","Ragged","batch","metadata",".","attention_mask","(","Optional","[","torch.Tensor","]",")",":","Attention","mask","tensor",".","If","None",",","masking","is","disabled",".","This","will","defer","to","the","config","in","the","case","of","conflicting","information",".","This","means","if","the","config","class","is","implying","causal","attention",",","the","mask","will","be","ignored",".","attention_bias","(","Optional","[","torch.Tensor","]",")",":","Attention","bias","tensor",".","If","None",",","bias","is","disabled","."]
code: def forward(
      self,
      q_k_v: torch.Tensor,
      kv_cache: torch.Tensor,
      batch: RaggedBatchWrapper,
      attention_mask: Optional[torch.Tensor] = None,
      attention_bias: Optional[torch.Tensor] = None,
      inv_freqs: Optional[torch.Tensor] = None,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Parameters:
          q_k_v (torch.Tensor): Query, key, and value tensors. Expected shape is:
              [
                  batch,
                  seq_len,
                  2 * self._config.n_heads_kv + self._config.n_heads_q,
                  self._config.head_size
              ].
          kv_cache (Optional[torch.Tensor]): Key and value cache tensor. Expected shape is
              [
                  2,
                  batch,
                  kv_cache_len,
                  self._config.n_heads_kv,
                  self._config.head_size
              ]. If None, cache is disabled. The `kv_cache_len` dimension does not need to
              be contiguous (it should expand stride by `max_out_tokens`).
          batch (RaggedBatchWrapper): Ragged batch metadata.
          attention_mask (Optional[torch.Tensor]): Attention mask tensor. If None, masking is
              disabled. This will defer to the config in the case of conflicting information.
              This means if the config class is implying causal attention, the mask will be ignored.
          attention_bias (Optional[torch.Tensor]): Attention bias tensor. If None, bias is disabled.
      """
      raise NotImplementedError()
  
file_path: DeepSpeed/deepspeed/inference/v2/modules/interfaces/attention_base.py
imports: from typing import Any, Dict, Optional, Tuple, Type
  import torch
  from ...ragged import RaggedBatchWrapper
  from deepspeed.runtime.config_utils import DeepSpeedConfigModel
  from ..ds_module import DSModuleBase
  from ..module_registry import DSModuleRegistryBase
  from ..configs import DSSelfAttentionConfig
label_desc: Use image node_id 6 for calling the DSSelfAttentionBase obj's underlying member method code with example usage: obj.forward(q_k_v, kv_cache, batch, attention_mask, attention_bias, inv_freqs) without return types