fix: 添加更多rope

Author: onenewcode

operators/src/attention/args.rs

@@ -53,9 +53,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/rope/args.rs

@@ -1,10 +1,44 @@
-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
+    Long {
+        long: ConstPtr<H>,
+        short: ConstPtr<H>,
+        max_pos: u32,
+        origin_pos: u32,
+    },
+    NtkParts {
+        alpha: f32,
+        beta: f32,
+        l0: f32,
+        s: f32,
+    },
+    Yarn {
+        alpha: f32,
+        beta: f32,
+        l0: f32,
+        s: f32,
+    },
+}
+
 pub struct Args<H: Hardware> {
     pub t_layout: TensorLayout,
     pub t_base: MutPtr<H>,
@@ -15,6 +49,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 {

operators/src/rope/common_cpu/mod.rs

@@ -1,11 +1,23 @@
-use super::{args::Meta, fill_pos, Args, Rope, Seq, SinCosTable};
+use std::ptr::null;
+
+use super::{args::Meta, args::RopeType as R, fill_pos, Args, Rope, Seq, SinCosTable};
 use crate::{
     common_cpu::Cpu, get_static, strides_not_support, ByteOf, LaunchError, QueueAlloc, SchemeError,
     Unsigned,
 };
 use digit_layout::{types as ty, DigitLayout};
 use half::f16;
-
+#[derive(Copy, Clone)]
+enum SchemeType<T> {
+    Rope,
+    Long {
+        long: *const T,
+        short: *const T,
+        s: f32,
+        origin_pos: u32,
+    },
+    Yarn,
+}
 pub struct Operator;
 
 impl Rope<Cpu> for Operator {
@@ -78,6 +90,7 @@ impl crate::Operator for Operator {
             p_layout,
             p_base,
             theta,
+            rope_type,
             ..
         } = args;
         let &[_, nh, dh] = t_layout.shape() else {
@@ -99,33 +112,89 @@ impl crate::Operator for Operator {
             return Err(strides_not_support("").into());
         }
 
-        macro_rules! calculate {
-            ($t:ty, $p:ty) => {
-                Scheme::<$t, $p> {
-                    nt,
-                    nh,
-                    dh,
-                    st,
-                    sh,
-                    sp,
-                    theta: *theta,
-                    t_base: t_base.cast(),
-                    p_base: p_base.cast(),
+        match rope_type {
+            R::Rope | R::Dyn { .. } | R::Ntk { .. } => {
+                let theta = match rope_type {
+                    R::Rope => *theta,
+                    R::Dyn { s, a } => theta * (a * s - a + 1.),
+                    R::Ntk { s } => theta * s,
+                    _ => unreachable!(),
+                };
+                macro_rules! calculate {
+                    ($t:ty, $p:ty) => {
+                        Scheme::<$t, $p> {
+                            nt,
+                            nh,
+                            dh,
+                            st,
+                            sh,
+                            sp,
+                            theta,
+                            t_base: t_base.cast(),
+                            p_base: p_base.cast(),
+                            scheme_type: SchemeType::Rope,
+                        }
+                        .calculate()
+                    };
                 }
-                .calculate()
-            };
-        }
 
-        use digit_layout::types as ty;
-        match (dt_t, dt_p) {
-            (ty::F16, ty::U32) => calculate!(f16, u32),
-            (ty::F16, ty::U64) => calculate!(f16, u64),
-            (ty::F32, ty::U32) => calculate!(f32, u32),
-            (ty::F32, ty::U64) => calculate!(f32, u64),
-            (ty::F64, ty::U32) => calculate!(f64, u32),
-            (ty::F64, ty::U64) => calculate!(f64, u64),
-            _ => todo!(),
+                use digit_layout::types as ty;
+                match (dt_t, dt_p) {
+                    (ty::F16, ty::U32) => calculate!(f16, u32),
+                    (ty::F16, ty::U64) => calculate!(f16, u64),
+                    (ty::F32, ty::U32) => calculate!(f32, u32),
+                    (ty::F32, ty::U64) => calculate!(f32, u64),
+                    (ty::F64, ty::U32) => calculate!(f64, u32),
+                    (ty::F64, ty::U64) => calculate!(f64, u64),
+                    _ => todo!(),
+                }
+            }
+            R::Long {
+                long,
+                short,
+                max_pos,
+                origin_pos,
+            } => {
+                let s = 1.0
+                    + ((*max_pos as f32 / *origin_pos as f32).ln() / (*origin_pos as f32).ln())
+                        .sqrt();
+                macro_rules! calculate {
+                    ($t:ty, $p:ty) => {
+                        Scheme::<$t, $p> {
+                            nt,
+                            nh,
+                            dh,
+                            st,
+                            sh,
+                            sp,
+                            theta: *theta,
+                            t_base: t_base.cast(),
+                            p_base: p_base.cast(),
+                            scheme_type: SchemeType::Long {
+                                long: long.cast(),
+                                short: short.cast(),
+                                s,
+                                origin_pos: *origin_pos,
+                            },
+                        }
+                        .calculate()
+                    };
+                }
+
+                use digit_layout::types as ty;
+                match (dt_t, dt_p) {
+                    (ty::F16, ty::U32) => calculate!(f16, u32),
+                    (ty::F16, ty::U64) => calculate!(f16, u64),
+                    (ty::F32, ty::U32) => calculate!(f32, u32),
+                    (ty::F32, ty::U64) => calculate!(f32, u64),
+                    (ty::F64, ty::U32) => calculate!(f64, u32),
+                    (ty::F64, ty::U64) => calculate!(f64, u64),
+                    _ => todo!(),
+                }
+            }
+            _ => {}
         }
+
         Ok(())
     }
 }
@@ -142,15 +211,15 @@ struct Scheme<A, P> {
     theta: f32,
     t_base: *mut A,
     p_base: *const P,
+    scheme_type: SchemeType<A>,
 }
 
 unsafe impl<A, P> Send for Scheme<A, P> {}
 unsafe impl<A, P> Sync for Scheme<A, P> {}
-
 /// 激活值。
 trait Activation: Sized {
     /// 激活值类型决定计算类型。
-    type Calculation;
+    type Calculation: Copy;
     /// 计算流程。
     fn calculate(pair: &mut [Self; 2], sin: Self::Calculation, cos: Self::Calculation);
 }
@@ -187,16 +256,37 @@ impl Activation for f64 {
 }
 
 trait Position<Calculation> {
-    fn freq_sin_cos(self, k: isize, dh: isize, theta: f32) -> (Calculation, Calculation);
+    fn freq_sin_cos_rope(self, k: isize, dh: isize, theta: f32) -> (Calculation, Calculation);
+    fn freq_sin_cos_long(
+        self,
+        k: isize,
+        dh: isize,
+        theta: f32,
+        factor: Calculation,
+        s: f32,
+    ) -> (Calculation, Calculation);
 }
 
 macro_rules! impl_position {
     ($a:ty) => {
         impl<T: Unsigned> Position<$a> for T {
             #[inline]
-            fn freq_sin_cos(self, k: isize, dh: isize, theta: f32) -> ($a, $a) {
+            fn freq_sin_cos_rope(self, k: isize, dh: isize, theta: f32) -> ($a, $a) {
                 (self.val() as $a / (theta as $a).powf(k as $a / dh as $a)).sin_cos()
             }
+            #[inline]
+            fn freq_sin_cos_long(
+                self,
+                k: isize,
+                dh: isize,
+                theta: f32,
+                factor: $a,
+                s: f32,
+            ) -> ($a, $a) {
+                let (sin, cos) =
+                    (self.val() as $a / (theta as $a).powf(k as $a / dh as $a) * factor).sin_cos();
+                (sin * s as $a, cos * s as $a)
+            }
         }
     };
 }
@@ -206,8 +296,8 @@ impl_position!(f64);
 
 impl<A, P> Scheme<A, P>
 where
-    A: Activation,
-    P: Position<A::Calculation> + Sync + Copy,
+    A: Activation + Copy,
+    P: Position<A::Calculation> + Sync + Copy + Unsigned,
 {
     fn calculate(&self) {
         let &Self {
@@ -220,6 +310,7 @@ where
             theta,
             t_base,
             p_base,
+            scheme_type,
         } = self;
         let nt = nt as isize;
         let nh = nh as isize;
@@ -229,10 +320,39 @@ where
         for i in 0..nt {
             let t = unsafe { t_base.byte_offset(i * st).cast::<[A; 2]>() };
             let p = unsafe { *p_base.byte_offset(i * sp) };
+            let factor = match scheme_type {
+                SchemeType::Long {
+                    long,
+                    short,
+                    origin_pos,
+                    ..
+                } => unsafe {
+                    if p.val() < origin_pos as usize {
+                        short.byte_offset(i * st).cast()
+                    } else {
+                        long.byte_offset(i * st).cast()
+                    }
+                },
+                _ => null(),
+            };
             for j in 0..nh {
                 for k in 0..dh {
                     let pair = unsafe { &mut *t.byte_offset(j * sh + k * sd) };
-                    let (sin, cos) = p.freq_sin_cos(k, dh, theta);
+                    let (sin, cos) = match scheme_type {
+                        SchemeType::Rope => p.freq_sin_cos_rope(k, dh, theta),
+                        SchemeType::Long {
+                            long,
+                            short,
+                            origin_pos,
+                            s,
+                        } => {
+                            let factor = unsafe { *factor };
+                            p.freq_sin_cos_long(k, dh, theta, factor, s)
+                        }
+                        _ => {
+                            todo!()
+                        }
+                    };
                     A::calculate(pair, sin, cos)
                 }
             }

operators/src/rope/cuda/mod.rs

@@ -184,7 +184,7 @@ extern "C" __global__ void {POS_U64}(
 #[cfg(test)]
 mod test {
     use super::{Args, Gpu, Operator, POS_U32, POS_U64};
-    use crate::{Hardware, Operator as _, TensorLayout};
+    use crate::{rope::args, Hardware, Operator as _, TensorLayout};
     use digit_layout::{
         types::{F16, F64, U32},
         DigitLayout,
@@ -203,6 +203,7 @@ mod test {
             cos_layout: TensorLayout::new_dyn(dt_t, &[dyn_(); 2], &[dyn_(); 2]),
             cos_base: null(),
             theta: 0.,
+            rope_type: args::RopeType::Rope,
         }
     }
 
@@ -227,6 +228,7 @@ mod test {
             cos_layout: TensorLayout::new_contiguous(dt_t, &[0, dh]),
             cos_base: null(),
             theta,
+            rope_type: args::RopeType::Rope,
         }
     }
 

operators/src/rope/mod.rs

@@ -8,7 +8,7 @@ pub mod infini;
 pub mod opencl;
 
 mod args;
-pub use args::Args;
+pub use args::{Args,RopeType};
 
 crate::op_trait! { Rope
     /// 生成 sincos 表（[2, n, dh]）。