code: def _fp_matmul(
      A,
      B,
      C,
      M,
      N,
      K,
      bias,
      stride_am,
      stride_ak,
      stride_bk,
      stride_bn,
      stride_cm,
      stride_cn,
      CACHE_M,
      CACHE_N,
      CACHE_K,
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
      BLOCK_K: tl.constexpr,
      GROUP_M: tl.constexpr,
      SPLIT_K: tl.constexpr,
      EVEN_K: tl.constexpr,
      ACC_TYPE: tl.constexpr,
      BIAS_ADD: tl.constexpr,
      ACTIVATION: tl.constexpr,
  ):
      # matrix multiplication
      pid = tl.program_id(0)
      pid_z = tl.program_id(1)
      grid_m = (M + BLOCK_M - 1) // BLOCK_M
      grid_n = (N + BLOCK_N - 1) // BLOCK_N
      # re-order program ID for better L2 performance
      width = GROUP_M * grid_n
      group_id = pid // width
      group_size = min(grid_m - group_id * GROUP_M, GROUP_M)
      pid_m = group_id * GROUP_M + (pid % group_size)
      pid_n = (pid % width) // (group_size)
      # do matrix multiplication
      rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
      rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
      ram = tl.max_contiguous(tl.multiple_of(rm % M, BLOCK_M), BLOCK_M)
      rbn = tl.max_contiguous(tl.multiple_of(rn % N, BLOCK_N), BLOCK_N)
      rk = pid_z * BLOCK_K + tl.arange(0, BLOCK_K)
      # pointers
      A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
      B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
      acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
      for k in range(K, 0, -BLOCK_K * SPLIT_K):
          if EVEN_K:
              a = tl.load(A)
              b = tl.load(B)
          else:
              a = tl.load(A, mask=rk[None, :] < k, other=0.0)
              b = tl.load(B, mask=rk[:, None] < k, other=0.0)
          acc += tl.dot(a, b)
          A += BLOCK_K * SPLIT_K * stride_ak
          B += BLOCK_K * SPLIT_K * stride_bk
      # bias addition
      if BIAS_ADD:
          bias_offset = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
          bias_ptr = bias + bias_offset
          b = tl.load(bias_ptr, mask=bias_offset < N)
          acc = acc + b[None, :]
      # activation
      if ACTIVATION == "relu":
          acc = tl.where(acc >= 0, acc, 0)
      elif ACTIVATION == "leaky_relu":
          acc = tl.where(acc >= 0, acc, 0.01 * acc)
      elif ACTIVATION == "gelu":
          # acc = tl.sigmoid(1.702 * acc) * acc
          acc = gelu_functor(acc)
      elif ACTIVATION == "sigmoid":
          acc = tl.sigmoid(acc)  # sigmoid
      acc = acc.to(C.dtype.element_ty)
      # rematerialize rm and rn to save registers
      rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
      rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
      C = C + (rm[:, None] * stride_cm + rn[None, :] * stride_cn)
      mask = (rm < M)[:, None] & (rn < N)[None, :]
      # handles write-back with reduction-splitting
      if SPLIT_K == 1:
          tl.store(C, acc, mask=mask)
      else:
          tl.atomic_add(C, acc, mask=mask)
  
file_path: DeepSpeed/deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py
imports: import triton
  import triton.language
  from .gelu import gelu_functor
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: _fp_matmul(A, B, C, M, N, K, bias, stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, CACHE_M, CACHE_N, CACHE_K, BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M, SPLIT_K, EVEN_K, ACC_TYPE, BIAS_ADD, ACTIVATION) without return types