flashinfer.rope.apply_rope_pos_ids_inplace¶

flashinfer.rope.apply_rope_pos_ids_inplace(q: Tensor, k: Tensor, pos_ids: Tensor, rotary_dim: int | None = None, interleave: bool = False, rope_scale: float = 1, rope_theta: float = 10000.0) → None¶

Apply rotary embedding to a batch of queries/keys (stored as RaggedTensor) in place. cos/sin values are computed on the fly inside the kernel.

Each element of pos_ids gives the position of the corresponding query/key row. Please see Ragged Tensor tutorial for more details about the ragged tensor.

Parameters:

q (torch.Tensor) – Query ragged tensor, shape: (nnz, num_q_heads, head_dim), where nnz is the number of query/key rows.
k (torch.Tensor) – Key ragged tensor, shape: (nnz, num_k_heads, head_dim).
pos_ids (torch.Tensor) – Position indices, shape: (nnz).
rotary_dim (Optional[int]) – The dimensions to which to apply RoPE. If None, RoPE is applied to the entire head dimension, otherwise it is applied to the first rotary_dim dimensions. Default: None.
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 rotate the first half dimensions ([..., :head_dim//2]) and the second half dimensions ([..., head_dim//2:]).
rope_scale (float) – The scaling factor used in the rope embedding, default: 1.
rope_theta (float) – The theta value used in the rope embedding, default: 1e4.