feat: 添加mmla的吸收版本

Author: onenewcode

operators/src/attention_mla/args.rs

@@ -0,0 +1,125 @@
+use crate::{
+    dyn_,
+    fuesd_softmax::AttnMask,
+    utils::{dim_distinct, rank_error, type_distinct},
+    ConstPtr, Hardware, MaybeDyn, MutPtr, SchemeError, TensorLayout,
+};
+use digit_layout::DigitLayout;
+use std::ptr::{null, null_mut};
+
+pub struct Args<H: Hardware> {
+    // q传入的是是吸收后的
+    pub q_layout: TensorLayout,
+    pub q_base: MutPtr<H>,
+
+    pub kv_layout: TensorLayout,
+    pub kv_base: ConstPtr<H>,
+
+    pub absorb_layout: TensorLayout,
+    pub absorb_base: ConstPtr<H>,
+
+    pub qr_layout: TensorLayout,
+    pub qr_base: ConstPtr<H>,
+
+    pub kr_layout: TensorLayout,
+    pub kr_base: ConstPtr<H>,
+
+    pub o_layout: TensorLayout,
+    pub o_base: MutPtr<H>,
+
+    pub mask: AttnMask,
+}
+
+pub(super) struct Meta {
+    pub dt: DigitLayout,
+    pub nh: MaybeDyn<usize>,
+    pub seq: MaybeDyn<usize>,
+    pub att: MaybeDyn<usize>,
+    pub dkv: MaybeDyn<usize>,
+    pub dv: MaybeDyn<usize>,
+    pub dr: MaybeDyn<usize>,
+}
+
+impl<H: Hardware> Args<H> {
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn new_null(
+        mask: AttnMask,
+        dt: DigitLayout,
+        nh: MaybeDyn<usize>,
+        dkv: MaybeDyn<usize>,
+        seq: MaybeDyn<usize>,
+        att: MaybeDyn<usize>,
+        dv: MaybeDyn<usize>,
+        dr: MaybeDyn<usize>,
+    ) -> Self {
+        let q_layout = TensorLayout::new_dyn(dt, &[nh, seq, dkv], &[dyn_(); 3]);
+        let kv_layout = TensorLayout::new_dyn(dt, &[nh, att, dkv], &[dyn_(); 3]);
+        let absorb_layout = TensorLayout::new_dyn(dt, &[nh, dv, dkv], &[dyn_(); 3]);
+        let qr_layout = TensorLayout::new_dyn(dt, &[nh, seq, dr], &[dyn_(); 3]);
+        let kr_layout = TensorLayout::new_dyn(dt, &[nh, att, dr], &[dyn_(); 3]);
+        let o_layout = TensorLayout::new_dyn(dt, &[nh, seq, dv], &[dyn_(); 3]);
+        Self {
+            q_layout,
+            q_base: null_mut(),
+            kv_layout,
+            kv_base: null(),
+            absorb_layout,
+            absorb_base: null(),
+            qr_layout,
+            qr_base: null(),
+            kr_layout,
+            kr_base: null(),
+            o_layout,
+            o_base: null_mut(),
+            mask,
+        }
+    }
+
+    pub(super) fn meta(&self) -> Result<Meta, SchemeError> {
+        let Self {
+            q_layout,
+            kv_layout,
+            absorb_layout,
+            qr_layout,
+            kr_layout,
+            o_layout,
+            ..
+        } = self;
+
+        let &[nh_q, seq_q, dkv_q] = q_layout.shape() else {
+            return Err(rank_error("q", 3, q_layout.ndim()));
+        };
+
+        let &[nh_kv, attn_kv, dkv_kv] = kv_layout.shape() else {
+            return Err(rank_error("kv", 3, kv_layout.ndim()));
+        };
+        let &[nh_a, dv_a, dkv_a] = absorb_layout.shape() else {
+            return Err(rank_error("absorb", 3, absorb_layout.ndim()));
+        };
+        let &[nh_qr, seq_qr, dr_qr] = qr_layout.shape() else {
+            return Err(rank_error("qr", 3, qr_layout.ndim()));
+        };
+        let &[nh_kr, att_kr, dr_kr] = kr_layout.shape() else {
+            return Err(rank_error("kr", 3, kr_layout.ndim()));
+        };
+        let &[nh_o, seq_o, dv_o] = o_layout.shape() else {
+            return Err(rank_error("o", 3, o_layout.ndim()));
+        };
+
+        Ok(Meta {
+            dt: type_distinct(&[
+                q_layout.dt(),
+                kv_layout.dt(),
+                qr_layout.dt(),
+                kr_layout.dt(),
+                o_layout.dt(),
+            ])?,
+            nh: dim_distinct(&[nh_q, nh_kv, nh_a, nh_qr, nh_kr, nh_o])?,
+            seq: dim_distinct(&[seq_q, seq_o, seq_qr])?,
+            att: dim_distinct(&[attn_kv, att_kr])?,
+            dkv: dim_distinct(&[dkv_a, dkv_kv, dkv_q])?,
+            dv: dim_distinct(&[dv_a, dv_o])?,
+            dr: dim_distinct(&[dr_kr, dr_qr])?,
+        })
+    }
+}

operators/src/attention_mla/common_cpu.rs

@@ -0,0 +1 @@
+impl_op!(common_cpu, Cpu);

operators/src/attention_mla/cuda.rs

@@ -0,0 +1 @@
+impl_op!(cuda, Gpu);

operators/src/attention_mla/infini.rs

@@ -0,0 +1 @@
+impl_op!(infini, Device);

operators/src/attention_mla/mod.rs

@@ -0,0 +1,26 @@
+mod args;
+mod operator;
+
+pub use args::Args;
+
+crate::op_trait!(AttentionMLA);
+
+macro_rules! impl_op {
+    ($dev:ident, $proc:ident) => {
+        pub type Operator = super::operator::Operator<
+            crate::$dev::$proc,
+            crate::mat_mul::$dev::Operator,
+            crate::fuesd_softmax::$dev::Operator,
+            crate::rearrange::$dev::Operator,
+        >;
+    };
+}
+
+#[cfg(any(use_cpu, test))]
+pub mod common_cpu;
+#[cfg(use_cuda)]
+pub mod cuda;
+#[cfg(use_infini)]
+pub mod infini;
+#[cfg(use_cl)]
+pub mod opencl;

operators/src/attention_mla/opencl.rs

@@ -0,0 +1 @@
+impl_op!(opencl, ClDevice);

operators/src/attention_mla/operator.rs

@@ -0,0 +1,203 @@
+use super::{args::Meta, Args, AttentionMLA};
+use crate::{
+    dyn_, fuesd_softmax, get_static, mat_mul, rearrange, ByteOf, Hardware, LaunchError, QueueAlloc,
+    SchemeError, TensorLayout, Workspace, WorkspaceCollector,
+};
+use ndarray_layout::ArrayLayout;
+use std::marker::PhantomData;
+
+pub struct Operator<Hardware, MatMul, Softmax, Rearrange> {
+    mat_mul: MatMul,
+    softmax: Softmax,
+    rearrange: Rearrange,
+    _phantom: PhantomData<Hardware>,
+}
+
+impl<H, M, S, R> AttentionMLA<H> for Operator<H, M, S, R>
+where
+    H: Hardware,
+    M: mat_mul::MatMul<H>,
+    S: fuesd_softmax::FusedSoftmax<H>,
+    R: rearrange::Rearrange<H>,
+{
+}
+
+impl<H, M, S, R> crate::Operator for Operator<H, M, S, R>
+where
+    H: Hardware,
+    M: mat_mul::MatMul<H>,
+    S: fuesd_softmax::FusedSoftmax<H>,
+    R: rearrange::Rearrange<H>,
+{
+    type Hardware = H;
+    type TopoNode = H;
+    type Args = Args<H>;
+
+    fn new(node: &Self::TopoNode) -> Self {
+        Self {
+            mat_mul: M::new(node),
+            softmax: S::new(node),
+            rearrange: R::new(node),
+            _phantom: PhantomData,
+        }
+    }
+
+    fn scheme(
+        &mut self,
+        args: &Self::Args,
+        max_workspace_size: usize,
+    ) -> Result<usize, SchemeError> {
+        // TODO
+        Ok(0)
+    }
+
+    fn launch<QA>(
+        &self,
+        args: &Self::Args,
+        workspace: &mut [ByteOf<Self::Hardware>],
+        queue_alloc: &QA,
+    ) -> Result<(), LaunchError>
+    where
+        QA: QueueAlloc<Hardware = Self::Hardware>,
+    {
+        let Meta {
+            dt,
+            nh,
+            seq,
+            att,
+            dkv,
+            dv,
+            dr,
+        } = args.meta()?;
+        let Args {
+            q_layout,
+            q_base,
+            kv_layout,
+            kv_base,
+            absorb_layout,
+            absorb_base,
+            qr_layout,
+            qr_base,
+            kr_layout,
+            kr_base,
+            o_layout,
+            o_base,
+            mask,
+        } = args;
+
+        let &[nh_skv, att_skv, dkv_skv] = kv_layout.strides() else {
+            unreachable!()
+        };
+        let &[nh_skr, att_skr, dr_skr] = kr_layout.strides() else {
+            unreachable!()
+        };
+        let &[nh_sa, dv_sa, dkv_sa] = absorb_layout.strides() else {
+            unreachable!()
+        };
+        let &[nh_so, seq_so, dv_so] = o_layout.strides() else {
+            unreachable!()
+        };
+        let ele = dt.nbytes();
+        get_static! {
+            nh      seq     dkv     dr
+            nh_skv  att_skv  dkv_skv
+            nh_skr  att_skr  dr_skr
+            nh_sa   dv_sa    dkv_sa
+            nh_so   seq_so   dv_so
+            dv      att
+        };
+
+        #[inline(always)]
+        fn layout(shape: [usize; 3], strides: [isize; 3]) -> ArrayLayout<3> {
+            ArrayLayout::new(&shape, &strides, 0)
+        }
+        let kv_first_layout = layout([nh, att, dkv], [nh_skv, att_skv, dkv_skv]).transpose(&[2, 1]);
+        let kr_layout = layout([nh, att, dr], [nh_skr, att_skr, dr_skr]).transpose(&[2, 1]);
+        let a_layout = layout([nh, dv, dkv], [nh_sa, dv_sa, dkv_sa]).transpose(&[2, 1]);
+        let att_w_layout = TensorLayout::new_contiguous(dt, &[nh, seq, att]);
+        let attn_t_layout = TensorLayout::new_contiguous(dt, &[nh, seq, dkv]);
+        let att_w_size = nh * seq * att * ele;
+        let att_t_size = nh * seq * dkv * ele;
+        let mut workspace = Workspace::new(queue_alloc, workspace, att_w_size + att_t_size);
+        let (att_w_buf, workspace) = workspace.split_at_mut(att_w_size);
+        let (attn_t_buf, workspace) = workspace.split_at_mut(att_t_size);
+
+        let kv_first_layout =
+            TensorLayout::new(dt, kv_first_layout.shape(), kv_first_layout.strides());
+        let kr_layout = TensorLayout::new(dt, kr_layout.shape(), kr_layout.strides());
+        let a_layout = TensorLayout::new(dt, a_layout.shape(), a_layout.strides());
+        // att_w = qr*kr^T + q*kv^T
+        self.mat_mul.launch(
+            &mat_mul::Args {
+                c_layout: att_w_layout.clone(),
+                c_base: att_w_buf.as_mut_ptr(),
+                beta: 0.,
+                a_layout: qr_layout.clone(),
+                a_base: *qr_base,
+                b_layout: kr_layout.clone(),
+                b_base: *kr_base,
+                alpha: ((dv + dr) as f32).sqrt().recip(),
+            },
+            workspace,
+            queue_alloc,
+        )?;
+       
+        self.mat_mul.launch(
+            &mat_mul::Args {
+                c_layout: att_w_layout.clone(),
+                c_base: att_w_buf.as_mut_ptr(),
+                beta: 1.,
+                a_layout: q_layout.clone(),
+                a_base: *q_base,
+                b_layout: kv_first_layout.clone(),
+                b_base: *kv_base,
+                alpha: ((dv + dr) as f32).sqrt().recip(),
+            },
+            workspace,
+            queue_alloc,
+        )?;
+        // att_w = softmax(att)
+        self.softmax.launch(
+            &fuesd_softmax::Args {
+                att_mask: *mask,
+                att_layout: att_w_layout.clone(),
+                att_base: att_w_buf.as_mut_ptr(),
+            },
+            workspace,
+            queue_alloc,
+        )?;
+        // attn_t=att_o*kv
+        self.mat_mul.launch(
+            &mat_mul::Args {
+                c_layout: attn_t_layout.clone(),
+                c_base: attn_t_buf.as_mut_ptr(),
+                beta: 0.,
+                a_layout: att_w_layout.clone(),
+                a_base: att_w_buf.as_ptr(),
+                b_layout: kv_layout.clone(),
+                b_base: *kv_base,
+                alpha: 1.,
+            },
+            workspace,
+            queue_alloc,
+        )?;
+
+        // attn =attn_t*absorb^T
+        self.mat_mul.launch(
+            &mat_mul::Args {
+                c_layout: o_layout.clone(),
+                c_base: *o_base,
+                beta: 0.,
+                a_layout: attn_t_layout.clone(),
+                a_base: attn_t_buf.as_ptr(),
+                b_layout: a_layout.clone(),
+                b_base: *absorb_base,
+                alpha: 1.,
+            },
+            workspace,
+            queue_alloc,
+        )?;
+
+        Ok(())
+    }
+}

operators/src/lib.rs

@@ -8,6 +8,7 @@ pub mod add_rows;
 pub mod all_reduce;
 pub mod attention;
 pub mod attention_kv_cached;
+pub mod attention_mla;
 pub mod broadcast;
 pub mod conv;
 pub mod fuesd_softmax;

operators/src/rope/common_cpu/mod.rs

@@ -1,12 +1,11 @@
-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;
+use std::ptr::null;
 #[derive(Copy, Clone)]
 enum NtkPartsType {
     None,