code: def test_flash_attn_qkvpacked(
      seqlen, d, dropout_p, causal, local, alibi, deterministic, dtype
  ):
      if (
          seqlen >= 2048
          and torch.cuda.get_device_properties("cuda").total_memory
          <= 16 * 2**30
      ):
          pytest.skip()  # Reference implementation OOM
      device = "cuda"
      # set seed
      torch.random.manual_seed(0)
      batch_size = 4
      nheads = 9
      window_size = (
          (-1, -1) if not local else torch.randint(0, seqlen, (2,))
      )
      qkv = torch.randn(
          batch_size,
          seqlen,
          3,
          nheads,
          d,
          device=device,
          dtype=dtype,
          requires_grad=True,
      )
      if alibi:
          alibi_slopes = (
              torch.rand(
                  batch_size, nheads, device=device, dtype=torch.float32
              )
              * 0.3
          )
          attn_bias = attn_bias_from_alibi_slopes(
              alibi_slopes, seqlen, seqlen, causal=causal
          )
      else:
          alibi_slopes, attn_bias = None, None
      out, lse, S_dmask = flash_attn_qkvpacked_func(
          qkv,
          dropout_p,
          causal=causal,
          window_size=window_size,
          alibi_slopes=alibi_slopes,
          deterministic=deterministic,
          return_attn_probs=True,
      )
      if dropout_p > 0.0:
          S_dmask_converted = convert_flash_attn_S_to_softmax(
              S_dmask,
              seqlen,
              seqlen,
              None,
              None,
              d,
              dropout_p > 0.0,
              causal=causal,
              window_size=window_size,
          )
          dropout_mask = S_dmask_converted >= 0
          attn_unnorm = S_dmask_converted.abs()
          attn = normalize_flash_attn_S(
              attn_unnorm,
              qkv[:, :, 0],
              qkv[:, :, 1],
              qkv[:, :, 2],
              None,
              None,
              attn_bias,
              dropout_p > 0.0,
              causal=causal,
              window_size=window_size,
          )
          dropout_fraction = get_dropout_fraction(
              dropout_mask,
              None,
              None,
              causal=causal,
              window_size=window_size,
          ).item()
          print(f"Actual dropout fraction: {dropout_fraction}")
      else:
          dropout_mask = None
  
      out_ref, attn_ref = attention_qkvpacked_ref(
          qkv,
          None,
          attn_bias,
          dropout_p,
          dropout_mask,
          causal=causal,
          window_size=window_size,
      )
      out_pt, attn_pt = attention_qkvpacked_ref(
          qkv,
          None,
          attn_bias,
          dropout_p,
          dropout_mask,
          causal=causal,
          window_size=window_size,
          upcast=False,
          reorder_ops=True,
      )
      # v = qkv[:, :, 2].float()
      # qk = torch.einsum('bshd,bthd->bhst', qkv[:, :, 0], qkv[:, :, 1]).float()
      # if causal:
      #     causal_mask = torch.triu(torch.ones(seqlen, seqlen, dtype=torch.bool, device=qkv.device), 1)
      #     qk.masked_fill_(causal_mask, float('-inf'))
      # m = qk.amax(-1, keepdim=True)
      # s_tmp = torch.exp((qk - m) / math.sqrt(d))
      # p_tmp = torch.softmax(qk / math.sqrt(d), -1)
      # p_dropped = p_tmp if dropout_mask is None else p_tmp.masked_fill(~dropout_mask, 0)
      # lse_ref = torch.logsumexp(qk / math.sqrt(d), -1)
      # qk_max1 = torch.max(qk[:, :, 128:, 192:], -1, keepdim=True).values
      # qk_max2 = torch.max(qk[:, :, 128:, 128:], -1, keepdim=True).values
      # qk_max3 = torch.max(qk[:, :, 128:, 64:], -1, keepdim=True).values
      # qk_max4 = torch.max(qk[:, :, 128:, :], -1, keepdim=True).values
      # o1 = torch.einsum('bhst,bthd->bshd', torch.exp((qk[:, :, 128:, 192:] - qk_max1) / math.sqrt(d)), v[:, 192:])
      # o2 = torch.einsum('bhst,bthd->bshd', torch.exp((qk[:, :, 128:, 128:] - qk_max2) / math.sqrt(d)), v[:, 128:])
      # o3 = torch.einsum('bhst,bthd->bshd', torch.exp((qk[:, :, 128:, 64:] - qk_max3) / math.sqrt(d)), v[:, 64:])
      # o4 = torch.einsum('bhst,bthd->bshd', torch.exp((qk[:, :, 128:, :] - qk_max4) / math.sqrt(d)), v[:, :])
      print(f"Output max diff: {(out - out_ref).abs().max().item()}")
      print(f"Output mean diff: {(out - out_ref).abs().mean().item()}")
      print(
          f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}"
      )
      print(
          f"Pytorch mean diff: {(out_pt - out_ref).abs().mean().item()}"
      )
      if dropout_p > 0.0:
          print(
              f"Attention max diff: {(attn - attn_ref).abs().max().item()}"
          )
          print(
              f"Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}"
          )
  
      g = torch.randn_like(out)
      # do_o = (g.float() * out.float()).sum(-1)
      # dv_tmp = torch.einsum('bhts,bthd->bshd', attn_pt[:, :, :64], g[:, :64])
      # dv_tmp1 = torch.einsum('bhts,bthd->bshd', attn_pt[:, :, 64:], g[:, 64:])
      if d <= MAX_HEADDIM_SM8x or (is_sm80 or is_sm90):
          (dqkv,) = torch.autograd.grad(out, qkv, g)
          (dqkv_ref,) = torch.autograd.grad(out_ref, qkv, g)
          (dqkv_pt,) = torch.autograd.grad(out_pt, qkv, g)
          print(
              f"dQ max diff: {(dqkv[:, :, 0] - dqkv_ref[:, :, 0]).abs().max().item()}"
          )
          print(
              f"dK max diff: {(dqkv[:, :, 1] - dqkv_ref[:, :, 1]).abs().max().item()}"
          )
          print(
              f"dV max diff: {(dqkv[:, :, 2] - dqkv_ref[:, :, 2]).abs().max().item()}"
          )
          print(
              f"dQKV mean diff: {(dqkv - dqkv_ref).abs().mean().item()}"
          )
          print(
              f"dQ Pytorch max diff: {(dqkv_pt[:, :, 0] - dqkv_ref[:, :, 0]).abs().max().item()}"
          )
          print(
              f"dK Pytorch max diff: {(dqkv_pt[:, :, 1] - dqkv_ref[:, :, 1]).abs().max().item()}"
          )
          print(
              f"dV Pytorch max diff: {(dqkv_pt[:, :, 2] - dqkv_ref[:, :, 2]).abs().max().item()}"
          )
          print(
              f"dQKV Pytorch mean diff: {(dqkv_pt - dqkv_ref).abs().mean().item()}"
          )
  
      # Check that FlashAttention's numerical error is at most twice the numerical error
      # of a Pytorch implementation.
      assert (out - out_ref).abs().max().item() <= 2 * (
          out_pt - out_ref
      ).abs().max().item()
  
      if dropout_p > 0.0:
          assert (attn - attn_ref).abs().max().item() <= 2 * (
              attn_pt - attn_ref
          ).abs().max().item()
          # With alibi, many of the prob values are 0.0 & -0.0 so dropout_fraction isn't accurate
          if not alibi:
              assert abs(dropout_fraction - dropout_p) <= (
                  0.01 if not local else 0.025
              )
  
      if d <= MAX_HEADDIM_SM8x or (is_sm80 or is_sm90):
          assert (dqkv - dqkv_ref).abs().max().item() <= 2 * (
              dqkv_pt - dqkv_ref
          ).abs().max().item()
  
imports: import math
  import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from flash_attn import flash_attn_func, flash_attn_kvpacked_func, flash_attn_qkvpacked_func, flash_attn_varlen_func, flash_attn_varlen_kvpacked_func, flash_attn_varlen_qkvpacked_func, flash_attn_with_kvcache
  from flash_attn.bert_padding import pad_input, unpad_input
  from flash_attn.flash_attn_interface import _get_block_size_n
  from flash_attn.layers.rotary import apply_rotary_emb
label_desc: Use image node_id 13 for calling a global function with example usage: test_flash_attn_qkvpacked(seqlen, d, dropout_p, causal, local, alibi, deterministic, dtype) without return types