code: def apply_rotary_emb_qkv_(
      qkv,
      cos,
      sin,
      cos_k=None,
      sin_k=None,
      interleaved=False,
      seqlen_offsets: Union[int, torch.Tensor] = 0,
  ):
      """
      Arguments:
          qkv: (batch_size, seqlen, 3, nheads, headdim)
          cos, sin: (seqlen, rotary_dim / 2)
          cos_k, sin_k: (seqlen, rotary_dim / 2), optional
          interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead of
              1st half and 2nd half (GPT-NeoX style).
          seqlen_offsets: (batch_size,) or int. Each sequence in Q and K is shifted by this amount.
              Most commonly used in inference when we have KV cache.
      Return:
          qkv: (batch_size, seqlen, 3, nheads, headdim)
      rotary_dim must be <= headdim
      Apply rotary embedding *inplace* to the first rotary_dim of Q and K.
      """
      return ApplyRotaryEmbQKV_.apply(
          qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets
      )
  
imports: import math
  from typing import Optional, Tuple, Union
  import torch
  from einops import rearrange, repeat
  from flash_attn.ops.triton.rotary import apply_rotary
label_desc: Use image node_id 4 for calling a global function with example usage: apply_rotary_emb_qkv_(qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets) and returns: ApplyRotaryEmbQKV_