code: def rope_padded(
      xq: torch.Tensor,
      xk: torch.Tensor,
      xv: torch.Tensor,
      cache_k: torch.Tensor,
      cache_v: torch.Tensor,
      attn_bias: BlockDiagonalCausalWithOffsetPaddedKeysMask,
      *,
      theta: float = 10000.0,
      out_q: Optional[torch.Tensor] = None,
      adjacents: bool = True,
      internal_dtype: str = "",
  ):
      """
      Performs RoPE (rotary embeddings) and kv-cache emplacement for a heterogeneous
      batch for inference in the style given by
      BlockDiagonalCausalWithOffsetPaddedKeysMask.
      The batch is concatenated along the sequence dimension, so the
      actual dim-0 length of all tensors is 1.
  
      xq, xk and xv should be (1, slen, n_heads, dim), where
      xq's n_heads can differ from xk and xv.
  
      This function places the roped xk in the right place in cache_k, and
      xv (unmodified) in the right place in cache_v, and returns out_q
      (the roped xq) such that things are ready to call
  
      xformers.ops.memory_efficient_attention(
          out_q, cache_k, cache_v, attn_bias=attn_bias
      )
  
      This functionality is experimental. Its API might be changed without warnings.
      Use it at your own risk.
  
      Arguments:
          xq: tensor of queries to apply rope to
          xk: tensor of keys to apply rope to
          xv: tensor of values to copy into cache_v
          cache_k: cache of keys, MODIFIED IN PLACE
          cache_v: cache of values, MODIFIED IN PLACE
          attn_bias: details the layout of caches.
                  Used to determine frequencies for the
                  RoPE calculation as well as the locations in cache_k and cache_v
                  to write to. Must be on the device.
          adjacents: If True, the inputs are in adjacent pairs along the final dim axis.
                    This is like the released LLaMA model.
                    If False, the dim axis is split in two equal pieces.
                     I.e. the features are ordered with all the real parts before all
                     the imaginary parts. This matches HuggingFace, e.g.
                     https://github.com/huggingface/transformers/blob/
                     f143037789288ba532dada934a118e648e715738/
                     src/transformers/models/llama/modeling_llama.py#L126-L130
          internal_dtype: set to "f32" or "f64" to enforce dtype in the calculation
      """
      if torch.is_grad_enabled() and (
          xq.requires_grad
          or xk.requires_grad
          or xv.requires_grad
          or cache_k.requires_grad
          or cache_v.requires_grad
          or out_q is not None
      ):
          raise ValueError("Gradients not supported.")
      assert _is_triton_available()
      import triton
  
      from ._triton.rope_padded_kernels import _rope_padded_kernel
  
      n_total_queries = attn_bias.q_seqinfo.seqstart_py[-1]
      cache_length = attn_bias.k_seqinfo.seqstart_py[-1]
      ndim = xq.ndim
      if ndim not in [4, 5]:
          raise ValueError("Unexpected xq dimension")
      xq_stride = xq.stride()
      xk_stride = xk.stride()
      xv_stride = xv.stride()
      cache_k_stride = cache_k.stride()
      cache_v_stride = cache_v.stride()
      cache_k_shape = cache_k.shape
      xk_shape = xk.shape
      n_kv_heads = xk_shape[-2]
      expected_kv_heads = n_kv_heads
      if xk_stride[-2] == 0:
          n_kv_heads = 1
      expected_cache_heads = n_kv_heads
      if n_kv_heads == 1 and cache_k_stride[-2] == 0:
          # If there's 1 kv head, don't care how expanded
          # cache_k is. User might expand before or after rope.
          expected_cache_heads = cache_k_shape[-2]
  
      if ndim == 4:
          bsz, q_len, n_q_heads, dim = xq.shape
          assert q_len == n_total_queries
          if xk_shape != (1, n_total_queries, expected_kv_heads, dim):
              raise ValueError("unexpected k shape")
          if xv.shape != (1, n_total_queries, expected_kv_heads, dim):
              raise ValueError("unexpected v shape")
          if cache_k_shape != (
              1,
              cache_length,
              expected_cache_heads,
              dim,
          ):
              raise ValueError("unexpected cache_k shape")
          if cache_v.shape != (
              1,
              cache_length,
              expected_cache_heads,
              dim,
          ):
              raise ValueError("unexpected cache_v shape")
          n_groups = 1
          out_q_stride: Tuple[int, ...] = (0, n_q_heads * dim, dim, 1)
  
      else:
          bsz, q_len, n_groups, n_q_heads, dim = xq.shape
          assert q_len == n_total_queries
          if xk_shape != (
              1,
              n_total_queries,
              n_groups,
              expected_kv_heads,
              dim,
          ):
              raise ValueError("unexpected k shape")
          if xv.shape != (
              1,
              n_total_queries,
              n_groups,
              expected_kv_heads,
              dim,
          ):
              raise ValueError("unexpected v shape")
          if cache_k_shape != (
              1,
              cache_length,
              n_groups,
              expected_cache_heads,
              dim,
          ):
              raise ValueError("unexpected cache_k shape")
          if cache_v.shape != (
              1,
              cache_length,
              n_groups,
              expected_cache_heads,
              dim,
          ):
              raise ValueError("unexpected cache_v shape")
          out_q_stride = (
              0,
              n_q_heads * dim * n_groups,
              n_q_heads * dim,
              dim,
              1,
          )
  
      if bsz != 1:
          raise ValueError(
              "Expected batch size dimension to be 1 as batches should be concatenated."
          )
      if xq_stride[-1] != 1:
          raise ValueError("Each q head must be contiguous")
      if xk_stride[-1] != 1:
          raise ValueError("Each k head must be contiguous")
      if xv_stride[-1] != 1:
          raise ValueError("Each v head must be contiguous")
      if cache_k_stride[-1] != 1:
          raise ValueError("Each cache_k head must be contiguous")
      if cache_v_stride[-1] != 1:
          raise ValueError("Each cache_v head must be contiguous")
      n_total_heads = n_q_heads + 2 * n_kv_heads
      v_start = n_total_heads - n_kv_heads
      k_start = n_q_heads
      if out_q is None:
          out_q = xq.new_empty(xq.shape)
      else:
          if out_q.shape != xq.shape:
              raise ValueError("Unexpected shape of out_q")
          out_q_stride = out_q.stride()
          if out_q_stride[-1] != 1:
              raise ValueError("Each out_q head must be contiguous")
  
      assert out_q is not None
  
      logical_bsz = len(attn_bias.q_seqinfo.seqstart_py) - 1
  
      # Less than 64KB per feature: enqueue fused kernel
      MAX_FUSED_SIZE = 65536 // xq.element_size()
      BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(dim))
      BLOCK_SIZE = max(BLOCK_SIZE, 128)
      BLOCK_SIZE = min(BLOCK_SIZE, 4096)
      # heuristics for number of warps
      num_warps = min(max(BLOCK_SIZE // 256, 1), 8)
      device = xq.device
      # Move these to the right device, like fmha does.
      attn_bias.k_seqinfo.to(device)
      attn_bias.q_seqinfo.to(device)
      seqstartq = attn_bias.q_seqinfo.seqstart
      seqstartk = attn_bias.k_seqinfo.seqstart
      seqlenk = attn_bias.k_seqinfo.seqlen
      assert internal_dtype in ["", "f32", "f64"]
      # experiment with the order of dims here.
      _rope_padded_kernel[
          (
              logical_bsz,
              attn_bias.q_seqinfo.max_seqlen,
              n_total_heads * n_groups,
          )
      ](
          xq,
          xk,
          xv,
          out_q,
          cache_k,
          cache_v,
          seqstartq,
          seqstartk,
          seqlenk,
          theta,
          k_start,
          v_start,
          n_groups,
          dim,
          xq_stride[1],
          xq_stride[2] if ndim == 5 else 0,
          xq_stride[-2],
          xk_stride[1],
          xk_stride[2] if ndim == 5 else 0,
          xk_stride[-2],
          xv_stride[1],
          xv_stride[2] if ndim == 5 else 0,
          xv_stride[-2],
          cache_k_stride[1],
          cache_k_stride[2] if ndim == 5 else 0,
          cache_k_stride[-2],
          cache_v_stride[1],
          cache_v_stride[2] if ndim == 5 else 0,
          cache_v_stride[-2],
          seqstartq.stride(0),
          seqstartk.stride(0),
          seqlenk.stride(0),
          out_q_stride[1],
          out_q_stride[2] if ndim == 5 else 0,
          out_q_stride[-2],
          internal_dtype,
          const_batch_strides=False,
          cache_padding_length=0,
          seqlenk_shift=0,
          BLOCK_SIZE=BLOCK_SIZE,
          adjacents=adjacents,
          num_warps=num_warps,
      )
      return out_q
  
imports: from typing import Optional, Tuple
  import torch
  from xformers.ops.fmha.attn_bias import BlockDiagonalCausalWithOffsetPaddedKeysMask
  from ..None import _is_triton_available
label_desc: Use image node_id 1 for calling a global function with example usage: rope_padded(xq, xk, xv, cache_k, cache_v, attn_bias) and returns: out_q