code: def test_flash_attn_race_condition(
      seqlen_q, seqlen_k, d, dropout_p, causal, dtype
  ):
      device = "cuda"
      # set seed
      torch.random.manual_seed(0)
      batch_size = 60  # Sometimes we need large batch size for the race conditions to trigger
      nheads = 4
      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,
      )
      torch.random.manual_seed(42)
      out0, lse0, _ = flash_attn_func(
          q, k, v, dropout_p, causal=causal, return_attn_probs=True
      )
      g = torch.randn_like(out0)
      if d <= MAX_HEADDIM_SM8x or (is_sm80 or is_sm90):
          (
              dq0,
              dk0,
              dv0,
          ) = torch.autograd.grad(out0, (q, k, v), g)
          # Numerical error if we just do any arithmetic on dq
          dq_atol = 2 * ((dq0 + 0.3 - 0.3) - dq0).abs().max().item()
  
      for i in range(250):
          torch.random.manual_seed(42)
          out, lse, _ = flash_attn_func(
              q, k, v, dropout_p, causal=causal, return_attn_probs=True
          )
          assert torch.equal(out, out0)
          assert torch.equal(lse, lse0)
  
          if d <= MAX_HEADDIM_SM8x or (is_sm80 or is_sm90):
              (
                  dq,
                  dk,
                  dv,
              ) = torch.autograd.grad(out, (q, k, v), g)
              dq_equal = torch.allclose(dq, dq0, atol=dq_atol)
              if not dq_equal:
                  print(
                      f"Iter {i}, {dq_atol = }, dQ max diff: {(dq - dq0).abs().max().item()}"
                  )
              assert torch.equal(dv, dv0)
              assert torch.equal(dk, dk0)
              assert dq_equal
  
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 21 for calling a global function with example usage: test_flash_attn_race_condition(seqlen_q, seqlen_k, d, dropout_p, causal, dtype) without return types