code: def batched_gemm_fusion(
      Q,
      K,
      V,
      Out,  #
      stride_qz,
      stride_qh,
      stride_qm,
      stride_qk,  #
      stride_kz,
      stride_kh,
      stride_kn,
      stride_kk,  #
      stride_vz,
      stride_vh,
      stride_vk,
      stride_vn,  #
      stride_oz,
      stride_oh,
      stride_om,
      stride_on,  #
      Z,
      NH,
      N_CTX,  #
      BLOCK_M: tl.constexpr,
      BLOCK_DMODEL: tl.constexpr,  #
      BLOCK_N: tl.constexpr,
  ):
      start_m = tl.program_id(0)
      off_hz = tl.program_id(1)
      acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
  
      q_tile_ptr = tl.make_block_ptr(
          base=Q,
          shape=(Z, NH, N_CTX, BLOCK_DMODEL),
          strides=(stride_qz, stride_qh, stride_qm, stride_qk),
          offsets=(off_hz // NH, off_hz % NH, start_m, 0),
          block_shape=(1, 1, BLOCK_M, BLOCK_DMODEL),
          order=(3, 2, 1, 0),
      )
      k_tile_ptr = tl.make_block_ptr(
          base=K,
          shape=(Z, NH, N_CTX, BLOCK_DMODEL),
          strides=(stride_kz, stride_kh, stride_kn, stride_kk),
          offsets=(off_hz // NH, off_hz % NH, 0, 0),
          block_shape=(1, 1, BLOCK_N, BLOCK_DMODEL),
          order=(3, 2, 1, 0),
      )
      v_tile_ptr = tl.make_block_ptr(
          base=V,
          shape=(Z, NH, N_CTX, BLOCK_DMODEL),
          strides=(stride_vz, stride_vh, stride_vk, stride_vn),
          offsets=(off_hz // NH, off_hz % NH, 0, 0),
          block_shape=(1, 1, BLOCK_N, BLOCK_DMODEL),
          order=(3, 2, 1, 0),
      )
      o_tile_ptr = tl.make_block_ptr(
          base=Out,
          shape=(Z, NH, N_CTX, BLOCK_DMODEL),
          strides=(stride_oz, stride_oh, stride_om, stride_on),
          offsets=(off_hz // NH, off_hz % NH, start_m, 0),
          block_shape=(1, 1, BLOCK_M, BLOCK_DMODEL),
          order=(3, 2, 1, 0),
      )
  
      q = tl.load(q_tile_ptr, boundary_check=(0, 1, 2, 3))
      q = tl.view(q, (BLOCK_M, BLOCK_DMODEL))
      for i in range(0, N_CTX, BLOCK_N):
          k = tl.load(k_tile_ptr, boundary_check=(0, 1, 2, 3))
          k = tl.view(k, (BLOCK_N, BLOCK_DMODEL))
          qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
          qk += tl.dot(q, tl.trans(k))
  
          p = qk.to(tl.float16)
          v = tl.load(v_tile_ptr, boundary_check=(0, 1, 2, 3))
          v = tl.view(v, (BLOCK_N, BLOCK_DMODEL))
          acc += tl.dot(p, v)
  
          k_tile_ptr = tl.advance(k_tile_ptr, [0, 0, BLOCK_N, 0])
          v_tile_ptr = tl.advance(v_tile_ptr, [0, 0, BLOCK_N, 0])
  
      acc = tl.view(acc, (1, 1, BLOCK_M, BLOCK_DMODEL))
      acc = acc.to(tl.float16)
      tl.store(o_tile_ptr, acc)
  
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 3 for calling a global function with example usage: batched_gemm_fusion(Q, K, V, Out, stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn, stride_oz, stride_oh, stride_om, stride_on, Z, NH, N_CTX, BLOCK_M, BLOCK_DMODEL, BLOCK_N) without return types