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.