code: def test_flash_attn_kvcache(
      seqlen_q,
      seqlen_k,
      d,
      has_batch_idx,
      rotary_fraction,
      rotary_interleaved,
      seqlen_new_eq_seqlen_q,
      causal,
      local,
      alibi,
      new_kv,
      mha_type,
      num_splits,
      dtype,
  ):
      if seqlen_q > seqlen_k and new_kv:
          pytest.skip()
      if not new_kv and rotary_fraction > 0.0:
          pytest.skip()
      device = "cuda"
      # set seed
      torch.random.manual_seed(0)
      batch_size = 2
      batch_size_cache = (
          batch_size if not has_batch_idx else batch_size * 2
      )
      nheads = 6
      # rotary_dim must be a multiple of 16, and must be <= d
      rotary_dim = math.floor(int(rotary_fraction * d) / 16) * 16
      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
      )
      seqlen_new = (
          seqlen_q
          if seqlen_new_eq_seqlen_q
          else torch.randint(1, seqlen_q + 1, (1,)).item()
      )
      if new_kv:
          k = torch.randn(
              batch_size,
              seqlen_new,
              nheads_k,
              d,
              device=device,
              dtype=dtype,
          )
          v = torch.randn(
              batch_size,
              seqlen_new,
              nheads_k,
              d,
              device=device,
              dtype=dtype,
          )
      else:
          k, v = None, None
      k_cache = torch.randn(
          batch_size_cache,
          seqlen_k,
          nheads_k,
          d,
          device=device,
          dtype=dtype,
      )
      v_cache = torch.randn(
          batch_size_cache,
          seqlen_k,
          nheads_k,
          d,
          device=device,
          dtype=dtype,
      )
      cache_seqlens = torch.randint(
          0,
          # If we don't use seqlen_q in the case of causal and rotary, cos/sin won't be long enough
          (
              seqlen_k
              - (
                  seqlen_q
                  if (causal or local) and rotary_dim > 1
                  else seqlen_new
              )
              + 1
          )
          if new_kv
          else (seqlen_k + 1),
          (batch_size,),
          dtype=torch.int32,
          device=device,
      )
      arange = rearrange(
          torch.arange(seqlen_k, device=device), "s -> 1 s"
      )
      cache_seqlens_expanded = rearrange(cache_seqlens, "b -> b 1")
      key_padding_mask = arange < cache_seqlens_expanded + (
          seqlen_new if new_kv else 0
      )
      if has_batch_idx:
          cache_batch_idx = torch.randperm(
              batch_size_cache, dtype=torch.int32, device=device
          )[:batch_size]
      else:
          cache_batch_idx = None
      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,
              None,
              key_padding_mask,
              causal=causal,
          )
      else:
          alibi_slopes, attn_bias = None, None
      # cache_seqlens = torch.tensor([64], dtype=torch.int32, device=device)
      if rotary_dim > 0:
          angle = (
              torch.rand(seqlen_k, rotary_dim // 2, device=device)
              * 2
              * math.pi
          )
          cos = torch.cos(angle).to(dtype=dtype)
          sin = torch.sin(angle).to(dtype=dtype)
          if causal or local:
              q_ro = apply_rotary_emb(
                  q,
                  cos,
                  sin,
                  seqlen_offsets=cache_seqlens,
                  interleaved=rotary_interleaved,
              )
          else:
              q_ro = rearrange(
                  apply_rotary_emb(
                      rearrange(q, "b s h d -> b 1 (s h) d"),
                      cos,
                      sin,
                      seqlen_offsets=cache_seqlens,
                      interleaved=rotary_interleaved,
                  ),
                  "b 1 (s h) d -> b s h d",
                  s=seqlen_q,
              )
          # q_ro = q
          k_ro = apply_rotary_emb(
              k,
              cos,
              sin,
              seqlen_offsets=cache_seqlens,
              interleaved=rotary_interleaved,
          )
      else:
          cos, sin = None, None
          q_ro, k_ro = q, k
      # k_cache[:, 64:] = -1
      k_cache_ref = (
          k_cache if not has_batch_idx else k_cache[cache_batch_idx]
      ).clone()
      v_cache_ref = (
          v_cache if not has_batch_idx else v_cache[cache_batch_idx]
      ).clone()
      if new_kv:
          update_mask = torch.logical_and(
              cache_seqlens_expanded <= arange,
              arange < cache_seqlens_expanded + seqlen_new,
          )
          k_cache_ref[update_mask] = rearrange(
              k_ro, "b s ... -> (b s) ..."
          )
          v_cache_ref[update_mask] = rearrange(
              v, "b s ... -> (b s) ..."
          )
      k_cache_rep = repeat(
          k_cache_ref, "b s h d -> b s (h g) d", g=nheads // nheads_k
      )
      v_cache_rep = repeat(
          v_cache_ref, "b s h d -> b s (h g) d", g=nheads // nheads_k
      )
      out = flash_attn_with_kvcache(
          q,
          k_cache,
          v_cache,
          k,
          v,
          cos,
          sin,
          cache_seqlens,
          cache_batch_idx,
          causal=causal,
          window_size=window_size,
          rotary_interleaved=rotary_interleaved,
          alibi_slopes=alibi_slopes,
          num_splits=num_splits,
      )
      # out = flash_attn_with_kvcache(
      #     q, k_cache, v_cache, cache_seqlens=cache_seqlens, causal=causal, window_size=window_size
      # )
      # out = flash_attn_with_kvcache(q, k_cache, v_cache, causal=causal, window_size=window_size)
      # qk = torch.einsum("bqhd,bkhd->bhqk", q, k_cache_ref)
      # m = qk.amax(-1, keepdim=True)
      # s_tmp = torch.exp((qk - m) / math.sqrt(d))
      # o1 = torch.einsum('bhst,bthd->bshd', s_tmp, v_cache_ref)
      # lse_ref = torch.logsumexp(qk / math.sqrt(d), -1)
      # probs = torch.softmax(qk, dim=-1)
      out_ref, _ = attention_ref(
          q_ro,
          k_cache_rep,
          v_cache_rep,
          None,
          key_padding_mask,
          attn_bias,
          0.0,
          None,
          causal=causal,
          window_size=window_size,
      )
      out_pt, _ = attention_ref(
          q_ro,
          k_cache_rep,
          v_cache_rep,
          None,
          key_padding_mask,
          attn_bias,
          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()}"
      )
  
      # Check that FlashAttention's numerical error is at most twice the numerical error
      # of a Pytorch implementation.
      if new_kv:
          k_cache_select = (
              k_cache if not has_batch_idx else k_cache[cache_batch_idx]
          )
          v_cache_select = (
              v_cache if not has_batch_idx else v_cache[cache_batch_idx]
          )
          assert torch.allclose(
              k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3
          )
          assert torch.equal(v_cache_select, v_cache_ref)
      mult = 3 if not alibi else 5
      assert (out - out_ref).abs().max().item() <= mult * (
          out_pt - out_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 20 for calling a global function with example usage: test_flash_attn_kvcache(seqlen_q, seqlen_k, d, has_batch_idx, rotary_fraction, rotary_interleaved, seqlen_new_eq_seqlen_q, causal, local, alibi, new_kv, mha_type, num_splits, dtype) without return types