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.

void nvte_fused_mla_kv_rope_mxfp8(

const NVTETensor kv,

const NVTETensor k_rope,

const int k_rope_stride,

const NVTETensor cos,

const NVTETensor sin,

NVTETensor key,

NVTETensor value,

cudaStream_t stream,

)

Split multi-latent attention’s key/value, apply RoPE to the key and quantize key and value to MXFP8.

Each head of kv holds 128 key dimensions without RoPE followed by 128 value dimensions. The key is those 128 dimensions followed by the 64 dimensions of k_rope with RoPE applied, the same for all heads. For token t, batch entry b and i < 32, with (x1, x2) = (k_rope[t, b, 2i], k_rope[t, b, 2i + 1]):

key[t, b, h, 128 + i]      = x1 * cos[t, i]      - x2 * sin[t, i]
key[t, b, h, 128 + 32 + i] = x2 * cos[t, 32 + i] + x1 * sin[t, 32 + i]

This is the RoPE of Megatron-LM’s MLA kernels (interleaved pairs in, the two halves out), with their roundings: with BF16 cos and sin, fma(x1, cos, -(x2 * sin)) and fma(x1, sin, x2 * cos) with every operation rounded to BF16; with FP32 cos and sin, the same in FP32, rounded once to BF16.

key and value are written as the MXFP8 quantization of their [s, b * h * d] views, row-wise (32 elements of a head) and column-wise (32 tokens), with E4M3 data and compact E8M0 scaling factors: the bytes the MXFP8 quantizer produces for the BF16 key and value.

Requires SM 10.0+, an even number of heads and a sequence length that is a multiple of 32.

Parameters:

• kv – [in] Key/value, [s, b, h, 256] in BF16.

• k_rope – [in] Key RoPE dimensions shared by the heads, s * b rows (token-major) of 64 BF16 values.

• k_rope_stride – [in] Elements from one row of k_rope to the next, a multiple of 8.

• cos – [in] Cosines, [max_s, 64] in BF16 or FP32; row t is used for token t.

• sin – [in] Sines, same shape and type as cos.

• key – [inout] MXFP8 key of shape [s, b * h * 192], with row-wise and column-wise data.

• value – [inout] MXFP8 value of shape [s, b * h * 128], with row-wise and column-wise data.

• stream – [in] CUDA stream used for the operation.