class_docstr: Calculate QK in a ring-exchange style
code: def forward(
      ctx, sub_q, sub_k, batch_size, num_attention_heads, sub_seq_length
  ):
      # save tensor for backward
      ctx.save_for_backward(sub_q, sub_k)
      ctx.sub_seq_length = sub_seq_length
  
      # create local segment of attention score
      attention_score = torch.empty(
          batch_size * num_attention_heads,
          sub_seq_length,
          sub_seq_length * gpc.get_world_size(ParallelMode.SEQUENCE),
          dtype=sub_q.dtype,
          device=get_current_device(),
      )
  
      # compute local QK^T
      part_a = torch.matmul(sub_q, sub_k.transpose(2, 1))
      local_rank = gpc.get_local_rank(ParallelMode.SEQUENCE)
      local_world_size = gpc.get_world_size(ParallelMode.SEQUENCE)
      start_idx = local_rank * sub_seq_length
      end_idx = (local_rank + 1) * sub_seq_length
      attention_score[:, :, start_idx:end_idx] = part_a
  
      # compute QK^T in ring-all-reduce style
      for i in range(local_world_size - 1):
          sub_k = ring_forward(sub_k, ParallelMode.SEQUENCE)
          start_idx, end_idx = _calc_incoming_device_range(
              i, local_rank, local_world_size, sub_seq_length
          )
          part_a = torch.matmul(sub_q, sub_k.transpose(2, 1))
          attention_score[:, :, start_idx:end_idx] = part_a
  
      return attention_score
  
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 RingQK obj's underlying member method code with example usage: obj.forward(ctx, sub_q, sub_k, batch_size, num_attention_heads, sub_seq_length) and returns: attention_score