fused_rope.h

Functions

void nvte_fused_rope_forward(const NVTETensor input, const NVTETensor cu_seqlens, const NVTETensor freqs, const NVTETensor start_positions, NVTETensor output, const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank, const int s, const int b, const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b, const int stride_h, const int stride_d, cudaStream_t stream)

Apply rotary positional embedding to the input tensor.

Parameters:

input – [in] Input tensor for fused rope.
cu_seqlens – [in] The cumulative sum of sequence lengths tensor. (Required for the thd format, empty tensor for other formats)
freqs – [in] The freqs tensor.
start_positions – [in] The beginning offsets for applying RoPE embeddings.
output – [out] Output tensor.
qkv_format – [in] QKV format.
interleaved – [in] Whether to use interleaved rotary position embedding.
cp_size – [in] Context parallel world size.
cp_rank – [in] Context parallel rank.
s – [in] Length of the s dimension of input.
b – [in] Length of the b dimension of input.
h – [in] Length of the h dimension of input.
d – [in] Length of the d dimension of input.
d2 – [in] Length of the d dimension of freqs.
stride_s_or_t – [in] Stride of the s (sbhd/bshd)/t (thd) dimension of input.
stride_b – [in] Stride of the b dimension of input. (0 for thd).
stride_h – [in] Stride of the h dimension of input.
stride_d – [in] Stride of the d dimension of input.
stream – [in] CUDA stream used for the operation.

void nvte_fused_rope_backward(const NVTETensor output_grads, const NVTETensor cu_seqlens, const NVTETensor freqs, const NVTETensor start_positions, NVTETensor input_grads, const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank, const int s, const int b, const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b, const int stride_h, const int stride_d, cudaStream_t stream)

Compute the backward of the fused rope.

Parameters:

output_grads – [in] Incoming gradient tensor for backward.
cu_seqlens – [in] The cumulative sum of sequence lengths tensor. (Required for the thd format, empty tensor for other formats)
freqs – [in] The freqs tensor.
start_positions – [in] The beginning offsets for applying RoPE embeddings.
input_grads – [out] Input gradient tensor to calculate.
qkv_format – [in] QKV format.
interleaved – [in] Whether to use interleaved rotary position embedding.
cp_size – [in] Context parallel world size.
cp_rank – [in] Context parallel rank.
s – [in] Length of the s dimension of output_grads.
b – [in] Length of the b dimension of output_grads.
h – [in] Length of the h dimension of output_grads.
d – [in] Length of the d dimension of output_grads.
d2 – [in] Length of the d dimension of freqs.
stride_s_or_t – [in] Stride of the s (sbhd/bshd)/t (thd) dimension of output_grads.
stride_b – [in] Stride of the b dimension of output_grads. (0 for thd).
stride_h – [in] Stride of the h dimension of output_grads.
stride_d – [in] Stride of the d dimension of output_grads.
stream – [in] CUDA stream used for the operation.

void nvte_fused_qkv_rope_forward(const NVTETensor qkv_input, const NVTETensor q_freqs, const NVTETensor k_freqs, const NVTETensor start_positions, NVTETensor q_out, NVTETensor k_out, NVTETensor v_out, const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank, const int s, const int b, const int h, const int d, const int d2, const int qkv_split_arg_list_0, const int qkv_split_arg_list_1, const int qkv_split_arg_list_2, cudaStream_t stream)

Apply rotary positional embedding to the combined QKV input tensor.

Parameters:

qkv_input – [in] Combined QKV input tensor for fused rope.
q_freqs – [in] The freqs tensor for Q.
k_freqs – [in] The freqs tensor for K.
start_positions – [in] The beginning offsets for applying RoPE embeddings.
q_out – [out] Output tensor for Q.
k_out – [out] Output tensor for K.
v_out – [out] Output tensor for V.
qkv_format – [in] QKV format.
interleaved – [in] Whether to use interleaved rotary position embedding.
cp_size – [in] Context parallel world size.
cp_rank – [in] Context parallel rank.
s – [in] Length of the s dimension of input.
b – [in] Length of the b dimension of input.
h – [in] Length of the h dimension of input.
d – [in] Length of the d dimension of input.
d2 – [in] Length of the d dimension of freqs.
qkv_split_arg_list_0 – [in] The hidden size for Q.
qkv_split_arg_list_1 – [in] The hidden size for K.
qkv_split_arg_list_2 – [in] The hidden size for V.
stream – [in] CUDA stream used for the operation.

void nvte_fused_qkv_rope_backward(const NVTETensor q_grad_out, const NVTETensor k_grad_out, const NVTETensor v_grad_out, const NVTETensor q_freqs, const NVTETensor k_freqs, NVTETensor qkv_grad_input, const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank, const int s, const int b, const int h, const int d, const int d2, const int qkv_split_arg_list_0, const int qkv_split_arg_list_1, const int qkv_split_arg_list_2, cudaStream_t stream)

Compute the backward of the fused qkv rope.

Parameters:

q_grad_out – [in] Incoming gradient tensor for Q.
k_grad_out – [in] Incoming gradient tensor for K.
v_grad_out – [in] Incoming gradient tensor for V.
q_freqs – [in] The freqs tensor for Q.
k_freqs – [in] The freqs tensor for K.
qkv_grad_input – [out] Input gradient tensor to calculate.
qkv_format – [in] QKV format.
interleaved – [in] Whether to use interleaved rotary position embedding.
cp_size – [in] Context parallel world size.
cp_rank – [in] Context parallel rank.
s – [in] Length of the s dimension of input.
b – [in] Length of the b dimension of input.
h – [in] Length of the h dimension of input.
d – [in] Length of the d dimension of input.
d2 – [in] Length of the d dimension of freqs.
qkv_split_arg_list_0 – [in] The hidden size for Q.
qkv_split_arg_list_1 – [in] The hidden size for K.
qkv_split_arg_list_2 – [in] The hidden size for V.
stream – [in] CUDA stream used for the operation.