class_docstr: CUDA Kernel implementation that will perform rotary position embeddings on the queries and keys
  before copying into a blocked KV cache.
docstr: Args:
      head_size: The size of the attention head.
      q_ratio: Ratio of q heads to kv heads (for GQA)
      dtype: Data type for the input/output. Supported values are torch.float16 and torch.bfloat16.
code: def __init__(
      self,
      head_size: int,
      n_q_heads: int,
      n_kv_heads: int,
      dtype: torch.dtype,
      rotary_dim: int,
      theta_base: float,
  ) -> None:
      """
      Args:
          head_size: The size of the attention head.
          q_ratio: Ratio of q heads to kv heads (for GQA)
          dtype: Data type for the input/output. Supported values are torch.float16 and torch.bfloat16.
      """
  
      q_ratio = n_q_heads // n_kv_heads
  
      if head_size not in BlockedRotaryEmbeddings.supported_head_sizes:
          raise ValueError(
              "Unsupported head size: {}, supported_head_sizes are {}".format(
                  head_size,
                  BlockedRotaryEmbeddings.supported_head_sizes,
              )
          )
  
      if q_ratio not in BlockedRotaryEmbeddings.supported_q_ratios:
          raise ValueError(
              "Unsupported q_ratio: {}, supported_q_ratios are {}".format(
                  q_ratio, BlockedRotaryEmbeddings.supported_q_ratios
              )
          )
  
      if not isinstance(dtype, DtypeEnum):
          dtype = DtypeEnum(dtype)
  
      if dtype not in BlockedRotaryEmbeddings.supported_dtypes:
          raise ValueError(
              "Unsupported data type: {}, supported_dtypes are {}".format(
                  dtype, BlockedRotaryEmbeddings.supported_dtypes
              )
          )
  
      inf_module = RaggedOpsBuilder().load()
      self.kernel = inf_module.kv_rotary_embeddings
      self.head_size = head_size
      self.n_q_heads = n_q_heads
      self.n_kv_heads = n_kv_heads
      self.rotary_dim = rotary_dim
      self.theta_base = theta_base
  
file_path: DeepSpeed/deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_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 1 to create a new BlockedRotaryEmbeddings object from inherited base classes: DSKernelBase with example: obj = BlockedRotaryEmbeddings(head_size, n_q_heads, n_kv_heads, dtype, rotary_dim, theta_base)