def METHOD_NAME(
    num_tokens: int,
    num_heads: int,
    head_size: int,
    max_position: int,
    rotary_dim: int,
    dtype: torch.dtype,
    base: int = 10000,
) -> None:
    positions = torch.randint(0, max_position, (num_tokens, ), device='cuda')
    query = torch.randn(num_tokens,
                        num_heads * head_size,
                        dtype=dtype,
                        device='cuda')
    key = torch.randn(num_tokens,
                      num_heads * head_size,
                      dtype=dtype,
                      device='cuda')
    # Create the rotary embedding.
    inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2) / rotary_dim))
    t = torch.arange(max_position).float()
    freqs = torch.einsum('i,j -> ij', t, inv_freq.float())
    cos = freqs.cos()
    sin = freqs.sin()
    cos_sin_cache = torch.cat((cos, sin), dim=-1)
    cos_sin_cache = cos_sin_cache.to(dtype=dtype, device='cuda')
    # Run the kernel. The kernel is in-place, so we need to clone the inputs.
    out_query = query.clone()
    out_key = key.clone()
    rotary_embedding_neox(
        positions,
        out_query,
        out_key,
        head_size,
        cos_sin_cache,
    )
    # Run the reference implementation.
    ref_rotary_embedding = RefRotaryEmbeddingNeox(
        dim=rotary_dim,
        max_position_embeddings=max_position,
        base=base,
    ).to(dtype=dtype, device='cuda')
    ref_query, ref_key = ref_rotary_embedding(
        positions,
        query.view(num_tokens, num_heads, head_size),
        key.view(num_tokens, num_heads, head_size),
    )
    ref_query = ref_query.view(num_tokens, num_heads * head_size)
    ref_key = ref_key.view(num_tokens, num_heads * head_size)
    # Compare the results.
    assert torch.allclose(out_query, ref_query, atol=1e-3, rtol=1e-5)
    assert torch.allclose(out_key, ref_key, atol=1e-3, rtol=1e-5)