code: def transform_gqa_megatron(
      src_param: torch.Tensor,
      head_size: int,
      n_q_heads: int,
      n_kv_heads: int,
  ) -> torch.Tensor:
      assert (
          src_param.shape[0] == (2 * n_kv_heads + n_q_heads) * head_size
      )
  
      head_ratio = n_q_heads // n_kv_heads
  
      # Reshape to get the groups as the leading dimension
      groups_leading_view = src_param.reshape(
          n_kv_heads, 2 + head_ratio, head_size, -1
      )
      q_heads = groups_leading_view[:, :head_ratio, :, :].reshape(
          -1, groups_leading_view.shape[-1]
      )
      k_heads = groups_leading_view[:, head_ratio, :, :].reshape(
          -1, groups_leading_view.shape[-1]
      )
      v_heads = groups_leading_view[:, head_ratio + 1, :, :].reshape(
          -1, groups_leading_view.shape[-1]
      )
      # Squeeze will remove extra dimension for bias
      return torch.cat([q_heads, k_heads, v_heads], dim=0).squeeze()
  
file_path: DeepSpeed/deepspeed/inference/v2/model_implementations/common_parameters/qkv_parameters.py
imports: import torch
  from ...model_implementations.parameter_base import ParameterBase
label_desc: Use image node_id 2 for calling a global function with example usage: transform_gqa_megatron(src_param, head_size, n_q_heads, n_kv_heads) and returns: torch