code: def test_flash_attn_deterministic(
      seqlen_q, seqlen_k, swap_sq_sk, d, causal, 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"
      # set seed
      torch.random.manual_seed(0)
      batch_size = 4
      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,
      )
      out = flash_attn_func(
          q,
          k,
          v,
          0.0,
          causal=causal,
          window_size=window_size,
          deterministic=True,
      )
  
      g = torch.randn_like(out)
      if d <= MAX_HEADDIM_SM8x or (is_sm80 or is_sm90):
          dq0, dk0, dv0 = torch.autograd.grad(
              out, (q, k, v), g, retain_graph=True
          )
          for _ in range(50):
              dq, dk, dv = torch.autograd.grad(
                  out, (q, k, v), g, retain_graph=True
              )
              assert torch.equal(dv, dv0)
              assert torch.equal(dk, dk0)
              assert torch.equal(dq, dq0)
  
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 25 for calling a global function with example usage: test_flash_attn_deterministic(seqlen_q, seqlen_k, swap_sq_sk, d, causal, local, dtype) without return types