code: def _load_multi_mqa(
      config,
      prefix: str,
      weights,
      bias: bool,
      head_size,
      num_heads,
      hidden_size,
  ):
      if any("c_attn" in k for k in weights.routing.keys()):
          slice_ = weights._get_slice(f"{prefix}.c_attn.weight")
          shape = slice_.get_shape()
          world_size = weights.process_group.size()
          rank = weights.process_group.rank()
          if config.transpose:
              block_size = (shape[1] - 2 * head_size) // world_size
              start = rank * block_size
              stop = (rank + 1) * block_size
              assert (shape[1] - 2 * head_size) % world_size == 0
              q_tensor = slice_[:, start:stop]
              kv_tensor = slice_[:, -2 * head_size :]
              weight = torch.cat([q_tensor, kv_tensor], dim=1).T
          else:
              block_size = (shape[0] - 2 * head_size) // world_size
              start = rank * block_size
              stop = (rank + 1) * block_size
              assert (shape[0] - 2 * head_size) % world_size == 0
              q_tensor = slice_[start:stop]
              kv_tensor = slice_[-2 * head_size :]
              weight = torch.cat([q_tensor, kv_tensor], dim=0)
          if bias:
              slice_ = weights._get_slice(f"{prefix}.c_attn.bias")
              shape = slice_.get_shape()
              block_size = (shape[0] - 2 * head_size) // world_size
              assert (shape[0] - 2 * head_size) % world_size == 0
              start = rank * block_size
              stop = (rank + 1) * block_size
              q_tensor = slice_[start:stop]
              kv_tensor = slice_[-2 * head_size :]
              bias = torch.cat([q_tensor, kv_tensor], dim=0)
      else:
          if config.transpose:
              w = [
                  weights.get_sharded(
                      f"{prefix}.q_attn.weight", dim=1
                  ).T,
                  weights.get_tensor(f"{prefix}.kv_attn.weight").T,
              ]
              weight = torch.cat(w, dim=0)
          else:
              w = [
                  weights.get_sharded(f"{prefix}.q_attn.weight", dim=0),
                  weights.get_tensor(f"{prefix}.kv_attn.weight"),
              ]
              weight = torch.cat(w, dim=1)
  
          if bias:
              b = [
                  weights.get_sharded(f"{prefix}.q_attn.bias", dim=0),
                  weights.get_tensor(f"{prefix}.kv_attn.bias"),
              ]
              bias = torch.cat(b, dim=0)
          else:
              bias = None
  
      weight = weight.to(dtype=weights.dtype).to(device=weights.device)
      assert list(weight.shape) == [
          (num_heads + 2) * head_size,
          hidden_size,
      ], f"{weight.shape} != {[(num_heads + 2) * head_size, hidden_size]}"
      if bias is not None:
          bias = bias.to(dtype=weights.dtype).to(device=weights.device)
          assert list(bias.shape) == [
              (num_heads + 2) * head_size
          ], f"{weight.shape} != {[(num_heads + 2) * head_size]}"
      return TensorParallelColumnLinear(
          get_linear(weight, bias, config.quantize)
      )
  
imports: import torch
  import torch.distributed
  from torch import nn
  from transformers.activations import ACT2FN
  from typing import Optional, List, Tuple
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.layers import TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelHead, TensorParallelEmbedding, FastLayerNorm, get_linear
label_desc: Use image node_id 3 for calling a global function with example usage: _load_multi_mqa(config, prefix, weights, bias, head_size, num_heads, hidden_size) and returns: TensorParallelColumnLinear