flashinfer.rope.apply_rope_pos_ids¶
- flashinfer.rope.apply_rope_pos_ids(q: Tensor, k: Tensor, pos_ids: Tensor, rotary_dim: int | None = None, interleave: bool = False, rope_scale: float = 1, rope_theta: float = 10000.0) Tuple[Tensor, Tensor]¶
将旋转嵌入应用于一批查询/键（存储为 RaggedTensor）。cos/sin 值在内核内部实时计算。
我们使用
indptr表示批处理中每个片段的起始指针，第 i 个片段的查询是q[indptr[i]:indptr[i+1]]，第 i 个片段的键是k[indptr[i]:indptr[i+1]]，indptr的第一个元素始终为 0，indptr的最后一个元素是批处理中的查询/键总数。有关 Ragged Tensor 的更多详细信息，请参阅 Ragged Tensor 教程。- 参数:
q (torch.Tensor) – 查询 Ragged Tensor，形状：
(nnz, num_q_heads, head_dim)`, 其中 ``nnz是indptr的最后一个元素。k (torch.Tensor) – 键 Ragged Tensor，形状：
(nnz, num_k_heads, head_dim)，其中nnz是indptr的最后一个元素。pos_ids (torch.Tensor) – 位置索引，形状：
(batch_size + 1)。rotary_dim (Optional[int]) – 应用 RoPE 的维度，如果为
None，则将 RoPE 应用于整个 head 维度，否则将 RoPE 应用于前rotary_dim个维度，默认值：None。interleave (bool) –
是否在最后一个维度中使用交错布局，默认值：
False。如果为
True，则查询/键张量的最后一个维度是交错的，即我们旋转偶数维度([..., ::2])和奇数维度([..., 1::2])。如果为
False，则查询/键张量的最后一个维度不交错，即我们旋转前一半维度([..., :head_dim//2])和后一半维度([..., head_dim//2:])。
rope_scale (float) – rope 嵌入中使用的缩放因子，默认值：
1。rope_theta (float) – rope 嵌入中使用的 theta 值，默认值：
1e4。
- 返回值:
q_rope (torch.Tensor) – 旋转后的查询张量，形状：
(nnz, num_q_heads, head_dim)。k_rope (torch.Tensor) – 旋转后的键张量，形状：
(nnz, num_k_heads, head_dim)。