libmir-cuda 0.2.0

CUDA inference backend for libmir
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

use crate::{
    PrefillDenseSwiGlu, PrefillMoeBlockBf16,
    backend::{block::PreparedDecodeMoeBlock, dense::graph::PreparedDecodeDense},
};

pub(in crate::backend::model) enum PreparedLayer {
    Moe(Box<PreparedDecodeMoeBlock>),
    Dense(Box<PreparedDecodeDense>),
}

impl PreparedLayer {
    pub const fn layer(&self) -> usize {
        match self {
            Self::Moe(prepared) => prepared.layer,
            Self::Dense(prepared) => prepared.block.attention.config.layer,
        }
    }
}

pub(in crate::backend::model) enum LayerPrefill<'a> {
    Moe(&'a mut PrefillMoeBlockBf16),
    Dense(&'a mut PrefillDenseSwiGlu),
}