fix: 添加更多rope

Author: onenewcode

operators/src/attention/args.rs

@@ -34,6 +34,7 @@ pub(super) struct Meta {
 }
 
 impl<H: Hardware> Args<H> {
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new_null(
         mask: AttnMask,
         dt: DigitLayout,
@@ -53,9 +54,9 @@ impl<H: Hardware> Args<H> {
             q_base: null_mut(),
             k_layout: k_layout.clone(),
             k_base: null(),
-            v_layout: v_layout,
+            v_layout,
             v_base: null(),
-            o_layout: o_layout,
+            o_layout,
             o_base: null_mut(),
             mask,
         }

operators/src/attention/operator.rs

@@ -65,7 +65,7 @@ where
         } = args;
 
         // 如果不能保证 nh seq att dh 已知，用任意值初始化算子
-        let (Some(&nh), Some(&seq), Some(&att), Some(&dh), Some(&dv)) = (
+        let (Some(&nh), Some(&seq), Some(&att), Some(&dh), Some(&_dv)) = (
             nh.get_static(),
             seq.get_static(),
             att.get_static(),
@@ -194,7 +194,7 @@ where
         let (att_buf, workspace) = workspace.split_at_mut(att_size);
 
         let head_group = nh / nkvh;
-        let (q_layout, qx_layout, q_base) = match qx {
+        let (_q_layout, qx_layout, q_base) = match qx {
             None => {
                 let q_layout = TensorLayout::new_contiguous(dt, &[nh, seq, dh]);
                 let qx_layout = TensorLayout::new_contiguous(dt, &[nkvh, head_group * seq, dh]);

operators/src/rearrange/args.rs

@@ -268,13 +268,13 @@ fn test_scheme() {
             dst_layout: TensorLayout::new(
                 F16,
                 &shape,
-                &[33554432 * 2, 16777216 * 2, 524288 * 2, 128 * 2, 1 * 2],
+                &[33554432 * 2, 16777216 * 2, 524288 * 2, 128 * 2, 2],
             ),
             dst_base: null_mut(),
             src_layout: TensorLayout::new(
                 F16,
                 &shape,
-                &[33554432 * 2, 16777216 * 2, 524288 * 2, 128 * 2, 1 * 2],
+                &[33554432 * 2, 16777216 * 2, 524288 * 2, 128 * 2, 2],
             ),
             src_base: null(),
         };

operators/src/rope/args.rs

@@ -1,10 +1,45 @@
-use crate::{
+use crate::{
     type_not_support,
     utils::{dim_distinct, rank_error},
     ConstPtr, Hardware, MaybeDyn, MutPtr, SchemeError, TensorLayout,
 };
 use digit_layout::DigitLayout;
 
+pub enum RopeType<H: Hardware> {
+    // 以下枚举通用一个 Scheme
+    Rope,
+    Pi {
+        s: f32,
+    },
+    Ntk {
+        s: f32,
+    },
+    Dyn {
+        s: f32,
+        a: f32,
+    },
+
+    // 以下枚举通用一个 Scheme
+    NtkParts {
+        alpha: f32,
+        beta: f32,
+        l0: f32,
+        s: f32,
+    },
+    Yarn {
+        alpha: f32,
+        beta: f32,
+        l0: f32,
+        s: f32,
+    },
+    Long {
+        long: ConstPtr<H>,
+        short: ConstPtr<H>,
+        max_pos: u32,
+        origin_pos: u32,
+    },
+}
+
 pub struct Args<H: Hardware> {
     pub t_layout: TensorLayout,
     pub t_base: MutPtr<H>,
@@ -15,6 +50,7 @@ pub struct Args<H: Hardware> {
     pub cos_layout: TensorLayout,
     pub cos_base: ConstPtr<H>,
     pub theta: f32,
+    pub rope_type: RopeType<H>,
 }
 
 pub(super) struct Meta {