class_docstr: Initialization.
  
  Arguments:
      local_attention (Module): local attention with q,k,v
      sequence_process_group (ProcessGroup): sequence parallel process group
      scatter_idx (int): scatter_idx for all2all comm
      gather_idx (int): gather_idx for all2all comm
docstr: forward
  
  Arguments:
      query (Tensor): query input to the layer
      key (Tensor): key input to the layer
      value (Tensor): value input to the layer
      args: other args
  
  Returns:
      * output (Tensor): context output
docstr_tok: ["forward","Arguments",":","query","(","Tensor",")",":","query","input","to","the","layer","key","(","Tensor",")",":","key","input","to","the","layer","value","(","Tensor",")",":","value","input","to","the","layer","args",":","other","args","Returns",":","*","output","(","Tensor",")",":","context","output"]
code: def forward(
      self, query: Tensor, key: Tensor, value: Tensor, *args: Any
  ) -> Tensor:
      """forward
  
      Arguments:
          query (Tensor): query input to the layer
          key (Tensor): key input to the layer
          value (Tensor): value input to the layer
          args: other args
  
      Returns:
          * output (Tensor): context output
      """
      # TODO Merge three alltoall calls into one
      # TODO (Reza): change the api on the megatron-deepspeed side so that we only receive all data (q,k, and v) together!
      # in shape : e.g.,  [s/p:h:]
      query_layer = _SeqAllToAll.apply(
          self.spg, query, self.scatter_idx, self.gather_idx
      )
      key_layer = _SeqAllToAll.apply(
          self.spg, key, self.scatter_idx, self.gather_idx
      )
      value_layer = _SeqAllToAll.apply(
          self.spg, value, self.scatter_idx, self.gather_idx
      )
  
      # out shape : e.g., [s:h/p:]
      context_layer = self.local_attn(
          query_layer, key_layer, value_layer, *args
      )
  
      output = _SeqAllToAll.apply(
          self.spg, context_layer, self.gather_idx, self.scatter_idx
      )
  
      # out e.g., [s/p::h]
      return output
  
file_path: deepspeed/deepspeed/sequence/layer.py
imports: import torch
  from typing import Any, Tuple
  from torch import Tensor
  from torch.nn import Module
  import deepspeed.comm
label_desc: Use image node_id 2 for calling the DistributedAttention obj's underlying member method code with example usage: obj.forward(query, key, value) and returns: output