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 {

operators/src/rope/common_cpu/mod.rs

@@ -1,11 +1,36 @@
-use super::{args::Meta, fill_pos, Args, Rope, Seq, SinCosTable};
+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;
+use std::ptr::null;
+#[derive(Copy, Clone)]
+enum NtkPartsType {
+    None,
+    Yarn,
+}
 
+#[derive(Copy, Clone)]
+enum SchemeType {
+    Rope {
+        s: f32,
+    },
+    Long {
+        long: *const u8,
+        short: *const u8,
+        s: f32,
+        origin_pos: u32,
+    },
+    NtkParts {
+        alpha: f32,
+        beta: f32,
+        l0: f32,
+        s: f32,
+        ntktype: NtkPartsType,
+    },
+}
 pub struct Operator;
 
 impl Rope<Cpu> for Operator {
@@ -78,6 +103,7 @@ impl crate::Operator for Operator {
             p_layout,
             p_base,
             theta,
+            rope_type,
             ..
         } = args;
         let &[_, nh, dh] = t_layout.shape() else {
@@ -99,6 +125,50 @@ impl crate::Operator for Operator {
             return Err(strides_not_support("").into());
         }
 
+        let (theta, scheme_type) = match rope_type {
+            R::Rope | R::Dyn { .. } | R::Ntk { .. } | R::Pi { .. } => {
+                let (theta, s) = match rope_type {
+                    R::Rope => (*theta, 1.),
+                    R::Dyn { s, a } => (theta * (a * s - a + 1.), 1.),
+                    R::Ntk { s } => (theta * s, 1.),
+                    R::Pi { s } => (*theta, *s),
+                    _ => unreachable!(),
+                };
+                (theta, SchemeType::Rope { s })
+            }
+            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();
+                let scheme_type = SchemeType::Long {
+                    long: long.cast(),
+                    short: short.cast(),
+                    s,
+                    origin_pos: *origin_pos,
+                };
+                (*theta, scheme_type)
+            }
+            R::Yarn { alpha, beta, l0, s } | R::NtkParts { alpha, beta, l0, s } => {
+                let ntktype = match rope_type {
+                    R::NtkParts { .. } => NtkPartsType::None,
+                    R::Yarn { .. } => NtkPartsType::Yarn,
+                    _ => unreachable!(),
+                };
+                let scheme_type = SchemeType::NtkParts {
+                    alpha: *alpha,
+                    beta: *beta,
+                    l0: *l0,
+                    s: *s,
+                    ntktype,
+                };
+                (*theta, scheme_type)
+            }
+        };
         macro_rules! calculate {
             ($t:ty, $p:ty) => {
                 Scheme::<$t, $p> {
@@ -108,9 +178,10 @@ impl crate::Operator for Operator {
                     st,
                     sh,
                     sp,
-                    theta: *theta,
+                    theta,
                     t_base: t_base.cast(),
                     p_base: p_base.cast(),
+                    scheme_type,
                 }
                 .calculate()
             };
@@ -142,15 +213,15 @@ struct Scheme<A, P> {
     theta: f32,
     t_base: *mut A,
     p_base: *const P,
+    scheme_type: SchemeType,
 }
 
 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,15 +258,69 @@ 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,
+        s: f32,
+    ) -> (Calculation, Calculation);
+    fn freq_sin_cos_long(
+        self,
+        k: isize,
+        dh: isize,
+        t: f32,
+        f: Calculation,
+        s: f32,
+    ) -> (Calculation, Calculation);
+    #[allow(clippy::too_many_arguments)]
+    fn freq_sin_cos_ntk_part(
+        self,
+        k: isize,
+        dh: isize,
+        theta: f32,
+        alpha: f32,
+        beta: f32,
+        l0: f32,
+        s: f32,
+        ntktype: NtkPartsType,
+    ) -> (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) {
-                (self.val() as $a / (theta as $a).powf(k as $a / dh as $a)).sin_cos()
+            fn freq_sin_cos_rope(self, k: isize, dh: isize, theta: f32, s: f32) -> ($a, $a) {
+                (self.val() as $a * s as $a * (theta as $a).powf(k as $a / dh as $a).recip())
+                    .sin_cos()
+            }
+            #[inline]
+            fn freq_sin_cos_long(self, k: isize, dh: isize, t: f32, f: $a, s: f32) -> ($a, $a) {
+                let (sin, cos) =
+                    (self.val() as $a * (t as $a).powf(k as $a / dh as $a).recip() * f).sin_cos();
+                (sin * s as $a, cos * s as $a)
+            }
+            #[inline]
+            fn freq_sin_cos_ntk_part(
+                self,
+                k: isize,
+                dh: isize,
+                theta: f32,
+                alpha: f32,
+                beta: f32,
+                l0: f32,
+                s: f32,
+                ntktype: NtkPartsType,
+            ) -> ($a, $a) {
+                use std::f32::consts::PI;
+                let pos = match ntktype {
+                    NtkPartsType::None => self.val() as $a,
+                    NtkPartsType::Yarn => self.val() as $a * (0.1 * s.ln() + 1.) as $a,
+                };
+                let theta = theta.powf(k as f32 / dh as f32).recip();
+                let r = ((l0 / (2. * PI / theta) - alpha) / (beta - alpha)).clamp(0., 1.);
+                (pos * ((1. - r) / s + r) as $a * theta as $a).sin_cos()
             }
         }
     };
@@ -206,8 +331,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 +345,7 @@ where
             theta,
             t_base,
             p_base,
+            scheme_type,
         } = self;
         let nt = nt as isize;
         let nh = nh as isize;
@@ -229,10 +355,38 @@ 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 as *const P).byte_offset(i * st).cast()
+                    } else {
+                        (long as *const P).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 { s } => p.freq_sin_cos_rope(k, dh, theta, s),
+                        SchemeType::Long { s, .. } => {
+                            let factor = unsafe { *factor };
+                            p.freq_sin_cos_long(k, dh, theta, factor, s)
+                        }
+                        SchemeType::NtkParts {
+                            alpha,
+                            beta,
+                            l0,
+                            s,
+                            ntktype,
+                        } => p.freq_sin_cos_ntk_part(k, dh, theta, alpha, beta, l0, s, ntktype),
+                    };
                     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,
         }
     }