code: def test_flash_attn_varlen_causal(
      seqlen_q, seqlen_k, swap_sq_sk, d, local, dtype
  ):
      if (
          max(seqlen_q, seqlen_k) >= 2048
          and torch.cuda.get_device_properties("cuda").total_memory
          <= 16 * 2**30
      ):
          pytest.skip()  # Reference implementation OOM
      if swap_sq_sk:
          seqlen_q, seqlen_k = seqlen_k, seqlen_q
      device = "cuda"
      causal = True
      # set seed
      torch.random.manual_seed(0)
      batch_size = 8
      nheads = 9
      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,
      )
      k = torch.randn(
          batch_size,
          seqlen_k,
          nheads,
          d,
          device=device,
          dtype=dtype,
          requires_grad=True,
      )
      v = torch.randn(
          batch_size,
          seqlen_k,
          nheads,
          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"
      )
      (
          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 = flash_attn_varlen_func(
          q_unpad,
          k_unpad,
          v_unpad,
          cu_seqlens_q,
          cu_seqlens_k,
          max_seqlen_q,
          max_seqlen_k,
          0.0,
          causal=causal,
          window_size=window_size,
      )
      out = output_pad_fn(out_unpad)
      out_ref, attn_ref = attention_ref(
          q,
          k,
          v,
          query_padding_mask,
          key_padding_mask,
          None,
          0.0,
          None,
          causal=causal,
          window_size=window_size,
      )
      out_pt, attn_pt = attention_ref(
          q,
          k,
          v,
          query_padding_mask,
          key_padding_mask,
          None,
          0.0,
          None,
          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()}"
      )
  
      g = torch.randn_like(out)
      do_o = (g.float() * out.float()).sum(-1)
      if d <= MAX_HEADDIM_SM8x or (is_sm80 or is_sm90):
          (
              dq_unpad,
              dk_unpad,
              dv_unpad,
          ) = torch.autograd.grad(out, (q_unpad, k_unpad, v_unpad), g)
          dq = dq_pad_fn(dq_unpad)
          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)
          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() + 1e-5
  
      if d <= MAX_HEADDIM_SM8x or (is_sm80 or is_sm90):
          assert (dq - dq_ref).abs().max().item() <= 2 * (
              dq_pt - dq_ref
          ).abs().max().item() + 1e-5
          assert (dk - dk_ref).abs().max().item() <= 2 * (
              dk_pt - dk_ref
          ).abs().max().item() + 1e-5
          assert (dv - dv_ref).abs().max().item() <= 2 * (
              dv_pt - dv_ref
          ).abs().max().item() + 1e-5
  
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 18 for calling a global function with example usage: test_flash_attn_varlen_causal(seqlen_q, seqlen_k, swap_sq_sk, d, local, dtype) without return types