code: def test_flash_attn_varlen_output(
      seqlen_q,
      seqlen_k,
      d,
      dropout_p,
      causal,
      local,
      alibi,
      deterministic,
      mha_type,
      dtype,
      kvpacked,
  ):
      if (
          max(seqlen_q, seqlen_k) >= 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
      nheads_k = (
          nheads
          if mha_type == "mha"
          else (1 if mha_type == "mqa" else 3)
      )
      assert nheads % nheads_k == 0
      window_size = (
          (-1, -1) if not local else torch.randint(0, seqlen_k, (2,))
      )
      q = torch.randn(
          batch_size,
          seqlen_q,
          nheads,
          d,
          device=device,
          dtype=dtype,
          requires_grad=True,
      )
      if kvpacked:
          kv = torch.randn(
              batch_size,
              seqlen_k,
              2,
              nheads_k,
              d,
              device=device,
              dtype=dtype,
              requires_grad=True,
          )
      else:
          k = torch.randn(
              batch_size,
              seqlen_k,
              nheads_k,
              d,
              device=device,
              dtype=dtype,
              requires_grad=True,
          )
          v = torch.randn(
              batch_size,
              seqlen_k,
              nheads_k,
              d,
              device=device,
              dtype=dtype,
              requires_grad=True,
          )
  
      query_padding_mask = generate_random_padding_mask(
          seqlen_q, batch_size, device, mode="random"
      )
      key_padding_mask = generate_random_padding_mask(
          seqlen_k, batch_size, device, mode="random"
      )
      # key_padding_mask = generate_random_padding_mask(seqlen_k, batch_size, device, mode='full')
      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_q,
              seqlen_k,
              query_padding_mask,
              key_padding_mask,
              causal=causal,
          )
      else:
          alibi_slopes, attn_bias = None, None
  
      if kvpacked:
          (
              q_unpad,
              kv_unpad,
              cu_seqlens_q,
              cu_seqlens_k,
              max_seqlen_q,
              max_seqlen_k,
              q,
              kv,
              output_pad_fn,
              dq_pad_fn,
              dkv_pad_fn,
          ) = generate_qkv(
              q,
              *kv.unbind(dim=2),
              query_padding_mask,
              key_padding_mask,
              kvpacked=True,
          )
          out_unpad, sm_lse, S_dmask = flash_attn_varlen_kvpacked_func(
              q_unpad,
              kv_unpad,
              cu_seqlens_q,
              cu_seqlens_k,
              max_seqlen_q,
              max_seqlen_k,
              dropout_p,
              causal=causal,
              window_size=window_size,
              alibi_slopes=alibi_slopes,
              deterministic=deterministic,
              return_attn_probs=True,
          )
      else:
          (
              q_unpad,
              k_unpad,
              v_unpad,
              cu_seqlens_q,
              cu_seqlens_k,
              max_seqlen_q,
              max_seqlen_k,
              q,
              k,
              v,
              output_pad_fn,
              dq_pad_fn,
              dk_pad_fn,
          ) = generate_qkv(
              q,
              k,
              v,
              query_padding_mask,
              key_padding_mask,
              kvpacked=False,
          )
          out_unpad, sm_lse, S_dmask = flash_attn_varlen_func(
              q_unpad,
              k_unpad,
              v_unpad,
              cu_seqlens_q,
              cu_seqlens_k,
              max_seqlen_q,
              max_seqlen_k,
              dropout_p,
              causal=causal,
              window_size=window_size,
              alibi_slopes=alibi_slopes,
              deterministic=deterministic,
              return_attn_probs=True,
          )
      out = output_pad_fn(out_unpad)
      if dropout_p > 0.0:
          S_dmask_converted = convert_flash_attn_S_to_softmax(
              S_dmask,
              seqlen_q,
              seqlen_k,
              query_padding_mask,
              key_padding_mask,
              d,
              dropout_p > 0.0,
              causal=causal,
              window_size=window_size,
          )
          dropout_mask = S_dmask_converted >= 0
          attn_unnorm = S_dmask_converted.abs()
          if kvpacked:
              kv_rep = repeat(
                  kv,
                  "b s two h d -> b s two (h g) d",
                  g=nheads // nheads_k,
              )
              k_rep, v_rep = kv_rep.unbind(dim=2)
          else:
              k_rep = repeat(
                  k, "b s h d -> b s (h g) d", g=nheads // nheads_k
              )
              v_rep = repeat(
                  v, "b s h d -> b s (h g) d", g=nheads // nheads_k
              )
          attn = normalize_flash_attn_S(
              attn_unnorm,
              q,
              k_rep,
              v_rep,
              query_padding_mask,
              key_padding_mask,
              attn_bias,
              dropout_p > 0.0,
              causal=causal,
              window_size=window_size,
          )
          dropout_fraction = get_dropout_fraction(
              dropout_mask,
              query_padding_mask,
              key_padding_mask,
              causal=causal,
              window_size=window_size,
          ).item()
          print(f"Actual dropout fraction: {dropout_fraction}")
      else:
          dropout_mask = None
  
      if kvpacked:
          out_ref, attn_ref = attention_kvpacked_ref(
              q,
              kv,
              query_padding_mask,
              key_padding_mask,
              attn_bias,
              dropout_p,
              dropout_mask,
              causal=causal,
              window_size=window_size,
          )
          out_pt, attn_pt = attention_kvpacked_ref(
              q,
              kv,
              query_padding_mask,
              key_padding_mask,
              attn_bias,
              dropout_p,
              dropout_mask,
              causal=causal,
              window_size=window_size,
              upcast=False,
              reorder_ops=True,
          )
      else:
          out_ref, attn_ref = attention_ref(
              q,
              k,
              v,
              query_padding_mask,
              key_padding_mask,
              attn_bias,
              dropout_p,
              dropout_mask,
              causal=causal,
              window_size=window_size,
          )
          out_pt, attn_pt = attention_ref(
              q,
              k,
              v,
              query_padding_mask,
              key_padding_mask,
              attn_bias,
              dropout_p,
              dropout_mask,
              causal=causal,
              window_size=window_size,
              upcast=False,
              reorder_ops=True,
          )
  
      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)
      if d <= MAX_HEADDIM_SM8x or (is_sm80 or is_sm90):
          if kvpacked:
              (
                  dq_unpad,
                  dkv_unpad,
              ) = torch.autograd.grad(out, (q_unpad, kv_unpad), g)
              dk, dv = dkv_pad_fn(dkv_unpad).unbind(2)
              (
                  dq_ref,
                  dkv_ref,
              ) = torch.autograd.grad(out_ref, (q, kv), g)
              dk_ref, dv_ref = dkv_ref.unbind(2)
              (
                  dq_pt,
                  dkv_pt,
              ) = torch.autograd.grad(out_pt, (q, kv), g)
              dk_pt, dv_pt = dkv_pt.unbind(2)
          else:
              (
                  dq_unpad,
                  dk_unpad,
                  dv_unpad,
              ) = torch.autograd.grad(
                  out, (q_unpad, k_unpad, v_unpad), g
              )
              dk = dk_pad_fn(dk_unpad)
              dv = dk_pad_fn(dv_unpad)
              (
                  dq_ref,
                  dk_ref,
                  dv_ref,
              ) = torch.autograd.grad(out_ref, (q, k, v), g)
              (
                  dq_pt,
                  dk_pt,
                  dv_pt,
              ) = torch.autograd.grad(out_pt, (q, k, v), g)
          dq = dq_pad_fn(dq_unpad)
          print(f"dQ max diff: {(dq - dq_ref).abs().max().item()}")
          print(f"dK max diff: {(dk - dk_ref).abs().max().item()}")
          print(f"dV max diff: {(dv - dv_ref).abs().max().item()}")
          print(f"dQ mean diff: {(dq - dq_ref).abs().mean().item()}")
          print(f"dK mean diff: {(dk - dk_ref).abs().mean().item()}")
          print(f"dV mean diff: {(dv - dv_ref).abs().mean().item()}")
          print(
              f"dQ Pytorch max diff: {(dq_pt - dq_ref).abs().max().item()}"
          )
          print(
              f"dK Pytorch max diff: {(dk_pt - dk_ref).abs().max().item()}"
          )
          print(
              f"dV Pytorch max diff: {(dv_pt - dv_ref).abs().max().item()}"
          )
          print(
              f"dQ Pytorch mean diff: {(dq_pt - dq_ref).abs().mean().item()}"
          )
          print(
              f"dK Pytorch mean diff: {(dk_pt - dk_ref).abs().mean().item()}"
          )
          print(
              f"dV Pytorch mean diff: {(dv_pt - dv_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 (dq - dq_ref).abs().max().item() <= 3 * (
              dq_pt - dq_ref
          ).abs().max().item()
          assert (dk - dk_ref).abs().max().item() <= 3 * (
              dk_pt - dk_ref
          ).abs().max().item()
          assert (dv - dv_ref).abs().max().item() <= 3 * (
              dv_pt - dv_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 16 for calling a global function with example usage: test_flash_attn_varlen_output(seqlen_q, seqlen_k, d, dropout_p, causal, local, alibi, deterministic, mha_type, dtype, kvpacked) without return types