code: def forward(ctx, q, k, v, sm_scale):
      BLOCK = 128
      # shape constraints
      Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
      assert Lq == Lk and Lk == Lv
      assert Lk in {16, 32, 64, 128}
      o = torch.empty_like(q)
      grid = (
          triton.cdiv(q.shape[2], BLOCK),
          q.shape[0] * q.shape[1],
          1,
      )
      L = torch.empty(
          (q.shape[0] * q.shape[1], q.shape[2]),
          device=q.device,
          dtype=torch.float32,
      )
      m = torch.empty(
          (q.shape[0] * q.shape[1], q.shape[2]),
          device=q.device,
          dtype=torch.float32,
      )
      num_warps = 4 if Lk <= 64 else 8
      D0 = q.shape[0] * q.shape[1] * q.shape[2]
      _fwd_kernel[grid](
          q,
          k,
          v,
          sm_scale,  #
          L,
          m,  #
          o,  #
          q.stride(0),
          q.stride(1),
          q.stride(2),
          q.stride(3),  #
          k.stride(0),
          k.stride(1),
          k.stride(2),
          k.stride(3),  #
          v.stride(0),
          v.stride(1),
          v.stride(2),
          v.stride(3),  #
          o.stride(0),
          o.stride(1),
          o.stride(2),
          o.stride(3),  #
          q.shape[0],
          q.shape[1],
          q.shape[2],
          D0,  #
          BLOCK_M=BLOCK,
          BLOCK_N=BLOCK,
          BLOCK_DMODEL=Lk,  #
          num_warps=num_warps,
          num_stages=2,
      )
  
      ctx.save_for_backward(q, k, v, o, L, m)
      ctx.grid = grid
      ctx.sm_scale = sm_scale
      ctx.BLOCK_DMODEL = Lk
      return o
  
file_path: triton/python/test/unit/hopper/test_flashattention.py
imports: import pytest
  import torch
  import triton
  import triton.language
label_desc: Use image node_id 1 for calling the _attention obj's underlying member method code with example usage: obj.forward(ctx, q, k, v, sm_scale) and returns: o