Instructions

Interface Definition

Initialize RopeLayerFWD. During construction, the maximum supported text context length, head dimension, and base for calculating the angle theta need to be passed.

RopeLayerFWD(const int maxPosition, const int dim, const float base = 10000,const float scalingFactor = 1.0f)->void

**Table 1** Input parameters of RopeLayerFWD
Parameter	Data Type	Description	Value Range
maxPosition	int	Maximum context length.	[1, INT_MAX]
dim	int	Head dimension.	[1, INT_MAX]
base	float	Base for calculating theta.	Base for calculating theta.
scalingFactor	float	Scaling factor of theta.	Scaling factor of theta.

Run the operator.

Run(void *query, void *key, const int qStride, const int kStride,const int batchSize, const int seqLen, const int qHeads, const int kHeads,const int *positionIds, KuDNN::RopeKind kind)->void

**Table 2** Input parameters of Run
Parameter	Data Type	Description	Value Range
query	void*	Pointer to the data to which RoPE is to be applied.	-
key	void*	Pointer to the data to which RoPE is to be applied.	-
qStride	int	Memory stride of the query.	[1, INT_MAX]
kStride	int	Memory stride of the key.	[1, INT_MAX]
batchSize	int	Batch size.	[1, INT_MAX]
seqLen	int	Statement length.	[1, INT_MAX]
qHeads	int	Number of query heads.	[1, INT_MAX]
kHeads	int	Number of key heads.	[1, INT_MAX]
positionlds	int*	List of token position IDs.	-
kind	KuDNN::RopeKind	(Optional) RoPE implementation type.	KuDNN::RopeKind::OpenSora KuDNN::RopeKind::DeepSeek

ValidateInput is used to validate the RopeLayerFWD input parameters and is automatically triggered during operator construction.

ValidateInput(const int maxPosition, const int dim, const float base, const float scalingFactor)->KuDNN::Status

**Table 3** Input parameters of ValidateInput
Parameter	Data Type	Description	Value Range
max_position	int	Maximum context length.	[1, INT_MAX]
dim	int	Head dimension.	[1, INT_MAX]
base	float	Base for calculating theta.	Base for calculating theta.
scaling_factor	float	Scaling factor of theta.	Scaling factor of theta.

Supported Data Types

RoPE supports the FP16 data type. (The query and key pointers point to the data type.)

**Table 4** Data types supported by query and key
query	key
fp16	fp16

The supported memory layout is (B, S, H, D).

B: batch size.
S: statement length, corresponding to seq_len.
H: number of heads, corresponding to head_num.
D: head dimension, corresponding to head_dim.

Examples

Perform RoPE on the query and key with max_position set to 10000 and head_dim set to 128.

int max_position = 10000, dim = 128;
int batchSize = 2, seqLen = 128, qHeads = 28, kHeads = 28; // The corresponding memory layout is (B, S, H, D).
// Construct the operator.
KuDNN::RopeLayerFWD layer(max_position, dim);
int qStride = qHeads * dim, kStride = kHeads * dim;
int qSize = batchSize * seqLen * qHeads * dim;
int kSize = batchSize * seqLen * kHeads * dim;
// Memory allocation for input and output parameters
__fp16 *query = (__fp16 *)malloc(qSize * sizeof(__fp16));
__fp16 *key = (__fp16 *)malloc(kSize * sizeof(__fp16));
int *positionIds = (int *)malloc(seqLen * sizeof(int));
// Run the operator.
layer.Run(query, key, qStride, kStride, batchSize, seqLen, qHeads, kHeads, positionIds);

Parent topic: RoPE