code: def group_gemm_fn(group_A, group_B):
      device = torch.device("cuda")
      assert len(group_A) == len(group_B)
      group_size = len(group_A)
  
      A_addrs = []
      B_addrs = []
      C_addrs = []
      g_sizes = []
      g_lds = []
      group_C = []
      for i in range(group_size):
          A = group_A[i]
          B = group_B[i]
          assert A.shape[1] == B.shape[0]
          M, K = A.shape
          K, N = B.shape
          C = torch.empty((M, N), device=device, dtype=A.dtype)
          group_C.append(C)
          A_addrs.append(A.data_ptr())
          B_addrs.append(B.data_ptr())
          C_addrs.append(C.data_ptr())
          g_sizes += [M, N, K]
          g_lds += [A.stride(0), B.stride(0), C.stride(0)]
  
      # note these are device tensors
      d_a_ptrs = torch.tensor(A_addrs, device=device)
      d_b_ptrs = torch.tensor(B_addrs, device=device)
      d_c_ptrs = torch.tensor(C_addrs, device=device)
      d_g_sizes = torch.tensor(
          g_sizes, dtype=torch.int32, device=device
      )
      d_g_lds = torch.tensor(g_lds, dtype=torch.int32, device=device)
      # we use a fixed number of CTA, and it's auto-tunable
      grid = lambda META: (META["NUM_SM"],)
      grouped_matmul_kernel[grid](
          d_a_ptrs,
          d_b_ptrs,
          d_c_ptrs,
          d_g_sizes,
          d_g_lds,
          group_size,
      )
  
      return group_C
  
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 2 for calling a global function with example usage: group_gemm_fn(group_A, group_B) and returns: group_C