class_docstr: CUDA Kernel implementation that will perform rotary position embeddings on the queries and keys
  before copying into a blocked KV cache.
docstr: Perform rotary embeddings on the queries and keys before copying into a blocked KV cache.
  
  Args:
      kv_cache (torch.Tensor): Pre-allocated KV cache of [num_blocks, block_size, 2, n_kv_heads, head_size]
      qkv: Input tensor of shape [num_tokens, head_size * (n_q_heads + 2 * n_kv_heads)]
      ragged_batch: Wrapper for the ragged batch.
      inverse_freqs: Inverse frequencies for the rotary embeddings. Shape [max_seq_len, rotary_dim // 2]
code: def __call__(
      self,
      kv_cache: torch.Tensor,
      qkv: torch.Tensor,
      ragged_batch: RaggedBatchWrapper,
      inverse_freqs: torch.Tensor,
  ) -> None:
      """
      Perform rotary embeddings on the queries and keys before copying into a blocked KV cache.
  
      Args:
          kv_cache (torch.Tensor): Pre-allocated KV cache of [num_blocks, block_size, 2, n_kv_heads, head_size]
          qkv: Input tensor of shape [num_tokens, head_size * (n_q_heads + 2 * n_kv_heads)]
          ragged_batch: Wrapper for the ragged batch.
          inverse_freqs: Inverse frequencies for the rotary embeddings. Shape [max_seq_len, rotary_dim // 2]
      """
  
      q = qkv[:, : self.head_size * self.n_q_heads]
      k = qkv[
          :,
          self.head_size
          * self.n_q_heads : self.head_size
          * (self.n_q_heads + self.n_kv_heads),
      ]
      v = qkv[:, self.head_size * (self.n_q_heads + self.n_kv_heads) :]
  
      self.kernel(
          kv_cache,
          q,
          k,
          v,
          inverse_freqs,
          ragged_batch.batch_metadata_buffer(),
          ragged_batch.inflight_seq_descriptors(),
          ragged_batch.tokens_to_seq(),
          ragged_batch.kv_ptrs(),
      )
  
file_path: DeepSpeed/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_trained_kv_rotary.py
imports: import torch
  from ....inference_utils import DtypeEnum
  from deepspeed.ops.op_builder import RaggedOpsBuilder
  from ....ragged import RaggedBatchWrapper
  from ...None import DSKernelBase
label_desc: Use image node_id 2 for calling the BlockedTrainedRotaryEmbeddings obj's underlying member method code with example usage: obj.__call__(kv_cache, qkv, ragged_batch, inverse_freqs) without return types