code: def forward(
      ctx,
      qkv,
      cos,
      sin,
      cos_k=None,
      sin_k=None,
      interleaved=False,
      seqlen_offsets: Union[int, torch.Tensor] = 0,
  ):
      batch, seqlen, three, nheads, headdim = qkv.shape
      assert three == 3
      if cos_k is None and sin_k is None and qkv.is_contiguous():
          # Call 1 kernel instead of 2 kernels
          # We need qkv to be contiguous so that when we reshape to combine (3, nheads)
          # dimensions, we get the same tensor
          # qk = rearrange(qkv[:, :, :2], "b s t h d -> b s (t h) d")
          qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim)
          apply_rotary(
              qk,
              cos,
              sin,
              seqlen_offsets=seqlen_offsets,
              interleaved=interleaved,
              inplace=True,
          )
      else:
          cos_k = cos if cos_k is None else cos_k
          sin_k = sin if sin_k is None else sin_k
          q, k = qkv[:, :, 0], qkv[:, :, 1]
          apply_rotary(
              q,
              cos,
              sin,
              seqlen_offsets,
              interleaved=interleaved,
              inplace=True,
          )
          apply_rotary(
              k,
              cos_k,
              sin_k,
              seqlen_offsets,
              interleaved=interleaved,
              inplace=True,
          )
          ctx.save_for_backward(cos, sin, cos_k, sin_k)
      if isinstance(seqlen_offsets, int):
          ctx.save_for_backward(cos, sin, cos_k, sin_k)
          ctx.seqlen_offsets = seqlen_offsets
      else:
          ctx.save_for_backward(cos, sin, cos_k, sin_k, seqlen_offsets)
          ctx.seqlen_offsets = None
      ctx.interleaved = interleaved
      return qkv
  
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 1 for calling the ApplyRotaryEmbQKV_ obj's underlying member method code with example usage: obj.forward(ctx, qkv, cos, sin, cos_k, sin_k, interleaved, seqlen_offsets) and returns: qkv