class_docstr: Modified implementation of flash-attn-2 tuned for inference on blocked KV-cache and wider
  range of input sequence lengths.
class_docstr_tok: ["Modified","implementation","of","flash-attn-2","tuned","for","inference","on","blocked","KV-cache","and","wider","range","of","input","sequence","lengths","."]
docstr: Flash attention implementation atop a blocked KV-cache. Atoms should be pre-populated.
  See attention_atom.h for further details on the structure of the information.
  
  Arguments:
      out (torch.Tensor): Output tensor of shape [tokens, hidden_size]
      q (torch.Tensor): Query tensor of shape [tokens, hidden_size]
      k (torch.Tensor): Key cache tensor of shape [n_blocks, block_size, n_heads_kv, head_size]. This Tensor only needs to be contiguous on the final dimension.
      v (torch.Tensor): Value cache tensor of shape [n_blocks, block_size, n_heads_kv, head_size]. This Tensor only needs to be contiguous on the final dimension.
      atoms (torch.Tensor): Atom information tensor of shape [num_atoms, 8] and type int32.
          Not all data is readable in this format. See attention_atom.h for further details.
      softmax_scale (float): Softmax scale factor.
  
  Returns:
      out (torch.Tensor): Output tensor of shape [tokens, hidden_size]
docstr_tok: ["Flash","attention","implementation","atop","a","blocked","KV-cache",".","Atoms","should","be","pre-populated",".","See","attention_atom.h","for","further","details","on","the","structure","of","the","information",".","Arguments",":","out","(","torch.Tensor",")",":","Output","tensor","of","shape","[","tokens",",","hidden_size","]","q","(","torch.Tensor",")",":","Query","tensor","of","shape","[","tokens",",","hidden_size","]","k","(","torch.Tensor",")",":","Key","cache","tensor","of","shape","[","n_blocks",",","block_size",",","n_heads_kv",",","head_size","]",".","This","Tensor","only","needs","to","be","contiguous","on","the","final","dimension",".","v","(","torch.Tensor",")",":","Value","cache","tensor","of","shape","[","n_blocks",",","block_size",",","n_heads_kv",",","head_size","]",".","This","Tensor","only","needs","to","be","contiguous","on","the","final","dimension",".","atoms","(","torch.Tensor",")",":","Atom","information","tensor","of","shape","[","num_atoms",",","8","]","and","type","int32",".","Not","all","data","is","readable","in","this","format",".","See","attention_atom.h","for","further","details",".","softmax_scale","(","float",")",":","Softmax","scale","factor",".","Returns",":","out","(","torch.Tensor",")",":","Output","tensor","of","shape","[","tokens",",","hidden_size","]"]
code: def __call__(
      self,
      out: torch.Tensor,
      q: torch.Tensor,
      k: torch.Tensor,
      v: torch.Tensor,
      atoms: torch.Tensor,
      softmax_scale: float,
  ) -> torch.Tensor:
      """
      Flash attention implementation atop a blocked KV-cache. Atoms should be pre-populated.
      See attention_atom.h for further details on the structure of the information.
  
      Arguments:
          out (torch.Tensor): Output tensor of shape [tokens, hidden_size]
          q (torch.Tensor): Query tensor of shape [tokens, hidden_size]
          k (torch.Tensor): Key cache tensor of shape [n_blocks, block_size, n_heads_kv, head_size]. This Tensor only needs to be contiguous on the final dimension.
          v (torch.Tensor): Value cache tensor of shape [n_blocks, block_size, n_heads_kv, head_size]. This Tensor only needs to be contiguous on the final dimension.
          atoms (torch.Tensor): Atom information tensor of shape [num_atoms, 8] and type int32.
              Not all data is readable in this format. See attention_atom.h for further details.
          softmax_scale (float): Softmax scale factor.
  
      Returns:
          out (torch.Tensor): Output tensor of shape [tokens, hidden_size]
      """
      self.kernel(out, q, k, v, atoms, softmax_scale, True)
      return out
  
file_path: deepspeed/deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.py
imports: import torch
  from deepspeed.accelerator import get_accelerator
  from ....inference_utils import DtypeEnum
  from deepspeed.ops.op_builder import RaggedOpsBuilder
  from ...None import DSKernelBase
label_desc: Use image node_id 2 for calling the BlockedFlashAttn obj's underlying member method code with example usage: obj.__call__(out, q, k, v, atoms, softmax_scale) and returns: out