class_docstr: fp16 matrix multiplication kernel
  dtypes: fp16 x fp16 = fp16
code: def _score_4d_matmul(input, head_size, input_mask, scale=-1.0):
      assert input.is_contiguous(), "matrix input must be contiguous"
  
      batches = input.shape[0]
      d_model = input.shape[-1] // 3
      num_of_heads = d_model // head_size
  
      q = input[:, :, :d_model]
      k = input[:, :, d_model : d_model * 2]
  
      q = q.view(batches, -1, num_of_heads, head_size)
      k = k.view(batches, -1, num_of_heads, head_size)
  
      # checks constraints
      assert q.shape == k.shape, "incompatible dimensions"
      B, M, H, K = q.shape
      B, N, H, K = k.shape
  
      assert K > 1, "inner-product dimension K should be larger than 1"
  
      # allocates output
      output = torch.empty((B, H, M, N), device=q.device, dtype=q.dtype)
      grid = lambda META: (
          triton.cdiv(M, META["BLOCK_SIZE_M"])
          * triton.cdiv(N, META["BLOCK_SIZE_N"]),
          H,
          B,
      )
      __class__._4d_kernel[grid](
          q,
          k,
          output,
          M,
          N,
          K,
          M // __class__._cache_stride,
          N // __class__._cache_stride,
          K // __class__._cache_stride,
          q.stride(0),
          q.stride(2),
          q.stride(1),
          q.stride(3),
          k.stride(0),
          k.stride(2),
          k.stride(3),
          k.stride(1),
          output.stride(0),
          output.stride(1),
          output.stride(2),
          output.stride(3),
          scale=scale,
          MASK=False,
      )
      return output
  
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 5 for calling the Fp16Matmul obj's underlying member method code with example usage: obj._score_4d_matmul(input, head_size, input_mask, scale) and returns: output