apply_rope_1d_blh

apply_rope_1d_blh(x, rope_1d_cache)

Apply 1D RoPE to a tensor laid out as [batch_size, seq_len, hidden_dim].

Parameters:

• x (Tensor) – Input tensor of shape [batch_size, seq_len, hidden_dim].

• rope_1d_cache (tuple[Tensor, Tensor]) – tuple[torch.Tensor, torch.Tensor] - The cache of 1D RoPE cos/sin for the input sequence.

Returns:

Tensor with the same shape as x.

Return type:

Tensor

Broadcasting:

• cos/sin are reshaped to [1, seq_len, hidden_dim].