code: def test_single_sequence_multiple_blocks(
      n_tokens: int,
      history_size: int,
      trained_emb: bool,
      head_size: int,
  ):
      """
      Validate that the copy works correctly
      """
      n_heads_q = 16
      n_heads_kv = 16
      kv_block_size = 64
      device = get_accelerator().current_device()
  
      batch, state_manager, seq_descs = build_batch_and_manager(
          [(n_tokens, history_size)],
          head_size,
          n_heads_kv,
          kv_block_size,
      )
  
      assert batch.current_sequences == 1
      assert batch.current_tokens == n_tokens
  
      qkv = torch.randn(
          (
              batch.current_tokens,
              (n_heads_q + 2 * n_heads_kv) * head_size,
          ),
          device=device,
          dtype=torch.float16,
      )
      qkv_ref = qkv.clone()
  
      q = qkv_ref[:, : head_size * n_heads_q]
      k = qkv_ref[
          :,
          head_size * n_heads_q : head_size * (n_heads_q + n_heads_kv),
      ]
      v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv) :]
  
      q_ref, k, freqs = rotary_pos_embs(
          q, k, seq_descs, batch, head_size
      )
      freqs = freqs.half()
  
      kv_cache = state_manager.get_cache(0)
  
      if trained_emb:
          copy_impl = BlockedTrainedRotaryEmbeddings(
              head_size, n_heads_q, n_heads_kv, torch.float16
          )
          copy_impl(kv_cache, qkv, batch, freqs)
      else:
          copy_impl = BlockedRotaryEmbeddings(
              head_size,
              n_heads_q,
              n_heads_kv,
              torch.float16,
              head_size,
              10000.0,
          )
          copy_impl(kv_cache, qkv, batch)
  
      assert allclose(qkv[:, : head_size * n_heads_q], q_ref)
      validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)
  
file_path: DeepSpeed/tests/unit/inference/v2/kernels/ragged_ops/test_blocked_rotary_emb.py
imports: from typing import List, Tuple
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.kernels.ragged_ops import BlockedRotaryEmbeddings, BlockedTrainedRotaryEmbeddings
  from deepspeed.inference.v2.ragged import RaggedBatchWrapper, DSSequenceDescriptor
  from .ragged_testing_utils import build_batch_and_manager, validate_kv_cache
  from ....v2.inference_test_utils import allclose
label_desc: Use image node_id 3 for calling a global function with example usage: test_single_sequence_multiple_blocks(n_tokens, history_size, trained_emb, head_size) without return types