class_docstr: Calculate AV in a ring-exchange style
code: def forward(
      ctx,
      attention_score,
      sub_v,
      batch_size,
      num_attention_heads,
      attention_head_size,
      sub_seq_length,
  ):
      local_rank = gpc.get_local_rank(ParallelMode.SEQUENCE)
      local_world_size = gpc.get_world_size(ParallelMode.SEQUENCE)
      local_start_idx, local_end_idx = _calc_current_device_range(
          local_rank, sub_seq_length
      )
  
      sub_attention_result = torch.zeros(
          batch_size * num_attention_heads,
          sub_seq_length,
          attention_head_size,
          device=get_current_device(),
          dtype=attention_score.dtype,
      )
  
      # save tensors for backward
      ctx.save_for_backward(attention_score, sub_v)
      ctx.sub_seq_length = sub_seq_length
  
      # compute local AV
      part_av = torch.matmul(
          attention_score[:, :, local_start_idx:local_end_idx], sub_v
      )
      sub_attention_result += part_av
  
      # compute AV in ring - all - reduce style
      for i in range(local_world_size - 1):
          sub_v = ring_forward(sub_v, ParallelMode.SEQUENCE)
          start_idx, end_idx = _calc_incoming_device_range(
              i, local_rank, local_world_size, sub_seq_length
          )
  
          # compute QK^T
          part_av = torch.matmul(
              attention_score[:, :, start_idx:end_idx], sub_v
          )
          sub_attention_result += part_av
      return sub_attention_result
  
file_path: ColossalAI/colossalai/legacy/nn/layer/parallel_sequence/_operation.py
imports: import torch
  from torch import distributed
  from torch.cuda.amp import custom_bwd, custom_fwd
  from colossalai.legacy.communication import ring_forward
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.nn.layer.parallel_sequence._utils import _calc_current_device_range, _calc_incoming_device_range
  from colossalai.utils import get_current_device
label_desc: Use image node_id 1 for calling the RingAV obj's underlying member method code with example usage: obj.forward(ctx, attention_score, sub_v, batch_size, num_attention_heads, attention_head_size, sub_seq_length) and returns: sub_attention_result