code: def test_batched_gemm_fusion():
      Z = 4
      NH = 48
      H = 64
      N_CTX = 2048
      BLOCK_M, BLOCK_N, BLOCK_DMODEL = 128, 128, H
      torch.manual_seed(20)
      A = torch.empty(
          (Z, NH, N_CTX, H), dtype=torch.float16, device="cuda"
      ).normal_(mean=0.1, std=0.2)
      B = torch.empty(
          (Z, NH, N_CTX, H), dtype=torch.float16, device="cuda"
      ).normal_(mean=0.1, std=0.2)
      C = torch.empty(
          (Z, NH, N_CTX, H), dtype=torch.float16, device="cuda"
      ).normal_(mean=0.1, std=0.2)
      E = torch.empty_like(A)
      BT = B.transpose(-1, -2)
      ref_out = torch.matmul(torch.matmul(A, BT), C)
      num_warps = 4
      grid = (triton.cdiv(N_CTX, BLOCK_M), B * NH)
      batched_gemm_fusion[grid](
          A,
          B,
          C,
          E,  #
          A.stride(0),
          A.stride(1),
          A.stride(2),
          A.stride(3),  #
          B.stride(0),
          B.stride(1),
          B.stride(2),
          B.stride(3),  #
          C.stride(0),
          C.stride(1),
          C.stride(2),
          C.stride(3),  #
          E.stride(0),
          E.stride(1),
          E.stride(2),
          E.stride(3),  #
          Z,
          NH,
          N_CTX,  #
          BLOCK_M,
          BLOCK_DMODEL,
          BLOCK_N,
          num_warps=num_warps,
      )
  
      torch.testing.assert_close(ref_out, E, atol=1e-2, rtol=0)
  
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 4 for calling a global function with example usage: test_batched_gemm_fusion() without return types