code: def flash_attn_with_kvcache(
      q,
      k_cache,
      v_cache,
      k=None,
      v=None,
      rotary_cos=None,
      rotary_sin=None,
      cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
      cache_batch_idx: Optional[torch.Tensor] = None,
      softmax_scale=None,
      causal=False,
      window_size=(-1, -1),  # -1 means infinite context window
      rotary_interleaved=True,
      num_splits=0,
  ):
      """
      If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
      k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
      the previous step, and update them with the new keys/values from the current step, and do
      attention with the updated cache, all in 1 kernel.
  
      If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
      For example, the KV cache could be pre-allocated with the max sequence length, and you can use
      cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.
  
      Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
      rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
      If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
      and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
      If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
      indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).
  
      See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.
  
      Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
      than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
      For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
      0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
  
      If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
      For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
          1 1 1 1 0
          1 1 1 1 1
      If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
          0 0
          0 0
          0 0
          1 0
          1 1
      If the row of the mask is all zero, the output will be zero.
  
      If window_size != (-1, -1), implements sliding window local attention. Query at position i
      will only attend to keys between
      [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
  
      Note: Does not support backward pass.
  
      Arguments:
          q: (batch_size, seqlen, nheads, headdim)
          k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim)
          v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim)
          k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
              k with k_cache, starting at the indices specified by cache_seqlens.
          v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.
          rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
              to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
          rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
          cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
              KV cache.
          cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
              If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
              If the indices are not distinct, and k and v are provided, the values updated in the cache
                   might come from any of the duplicate indices.
          softmax_scale: float. The scaling of QK^T before applying softmax.
              Default to 1 / sqrt(headdim).
          causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
          window_size: (left, right). If not (-1, -1), implements sliding window local attention.
          rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
              If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
              rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
              (i.e. GPT-NeoX style).
          num_splits: int. If > 1, split the key/value into this many chunks along the sequence.
             If num_splits == 1, we don't split the key/value. If num_splits == 0, we use a heuristic
             to automatically determine the number of splits.
             Don't change this unless you know what you are doing.
  
      Return:
          out: (batch_size, seqlen, nheads, headdim).
      """
      assert (
          k_cache.stride(-1) == 1
      ), "k_cache must have contiguous last dimension"
      assert (
          v_cache.stride(-1) == 1
      ), "v_cache must have contiguous last dimension"
      maybe_contiguous = (
          lambda x: x.contiguous()
          if x is not None and x.stride(-1) != 1
          else x
      )
      q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
      if softmax_scale is None:
          softmax_scale = q.shape[-1] ** (-0.5)
      if cache_seqlens is not None and isinstance(cache_seqlens, int):
          cache_seqlens = torch.full(
              (k_cache.shape[0],),
              cache_seqlens,
              dtype=torch.int32,
              device=k_cache.device,
          )
          cache_seqlens = maybe_contiguous(cache_seqlens)
      cache_batch_idx = maybe_contiguous(cache_batch_idx)
      out, softmax_lse = flash_attn_cuda.fwd_kvcache(
          q,
          k_cache,
          v_cache,
          k,
          v,
          cache_seqlens,
          rotary_cos,
          rotary_sin,
          cache_batch_idx,
          None,
          softmax_scale,
          causal,
          window_size[0],
          window_size[1],
          rotary_interleaved,
          num_splits,
      )
      return out
  
imports: from typing import Optional, Union
  import torch
  import torch.nn
  import flash_attn_2_cuda
label_desc: Use image node_id 12 for calling a global function with example usage: flash_attn_with_kvcache(q, k_cache, v_cache, k, v, rotary_cos, rotary_sin, cache_seqlens, cache_batch_idx, softmax_scale, causal, window_size, rotary_interleaved, num_splits) and returns: out