flashinfer.rope.apply_rope_with_cos_sin_cache_inplace#

flashinfer.rope.apply_rope_with_cos_sin_cache_inplace(q: torch.Tensor, k: torch.Tensor, cos_cache: torch.Tensor, sin_cache: torch.Tensor, pos_ids: torch.Tensor, interleave: bool = False) None#

Apply rotary embedding to keys and queries with precomputed cos/sin values. The result is stored in the input tensors inplace.

Parameters:
  • q (torch.Tensor) – Query tensor, shape: (nnz, num_q_heads, head_dim).

  • k (torch.Tensor) – Key tensor, shape: (nnz, num_k_heads, head_dim).

  • cos_cache (torch.Tensor) – Cosine cache tensor, shape: (max_seq_len, rotary_dim). Expect float32 data type.

  • sin_cache (torch.Tensor) – Sine cache tensor, shape: (max_seq_len, rotary_dim). Expect float32 data type.

  • pos_ids (torch.Tensor) – Position indices, shape: (nnz).

  • interleave (bool) –

    Whether to use interleaved layout in the last dimension, default: False.

    • If True, the last dimension of the query/key tensor is interleaved, i.e., we rotate the even dimensions ([..., ::2]) and odd dimensions ([..., 1::2]).

    • If False, the last dimension of the query/key tensor is not interleaved, i.e., we rorate the first half dimensions ([..., :head_dim//2]) and the second half dimensions ([..., head_dim//2:]).

Note

The rotary dimension is determined by the cosine cache and sine cache.