class_docstr: fp16 matrix multiplication kernel
  dtypes: fp16 x fp16 = fp16
code: def forward(A, B, use_triton=True, bias=None, activation=""):
      if use_triton:
          device = A.device
          # handle non-contiguous inputs if necessary
          if A.stride(0) > 1 and A.stride(1) > 1:
              A = A.contiguous()
          if B.stride(0) > 1 and B.stride(1) > 1:
              B = B.contiguous()
          # checks constraints
          assert A.shape[1] == B.shape[0], "incompatible dimensions"
          M, K = A.shape
          _, N = B.shape
          # allocates output
          C = torch.empty((M, N), device=device, dtype=A.dtype)
          # accumulator types
          ACC_TYPE = (
              triton.language.float32
              if A.dtype
              in [torch.float16, torch.bfloat16, torch.float32]
              else triton.language.int32
          )
          # launch kernel
          grid = lambda META: (
              triton.cdiv(M, META["BLOCK_M"])
              * triton.cdiv(N, META["BLOCK_N"]),
              META["SPLIT_K"],
          )
          __class__._2d_kernel[grid](
              A,
              B,
              C,
              M,
              N,
              K,
              bias,
              A.stride(0),
              A.stride(1),
              B.stride(0),
              B.stride(1),
              C.stride(0),
              C.stride(1),
              M // __class__._cache_stride,
              N // __class__._cache_stride,
              K // __class__._cache_stride,
              GROUP_M=8,
              ACC_TYPE=ACC_TYPE,
              BIAS_ADD=(0 if bias is None else 1),
              ACTIVATION=activation,
          )
      else:
          C = torch.matmul(A, B)
      return C
  
imports: import torch
  import triton
  import os
  from filelock import FileLock
  import deepspeed.ops.transformer.inference.triton.triton_matmul_kernel
  import pickle
  from io import open
  import deepspeed
  from pathlib import Path
  import atexit
label_desc: Use image node_id 3 for calling the Fp16Matmul obj's underlying member method code with example usage: obj.forward(A, B, use_triton, bias, activation) and returns: C