libmir-cuda 0.1.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
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

use super::{
    Bf16LinearPairWeights, CudaTensor, DecodeAttentionConfig, DecodeAttentionWeights,
    GatedActivation, RouterTensors,
};
use crate::{Error, Result};

#[derive(Clone, Copy, Debug)]
pub struct DecodeMoeBlockConfig {
    pub attention: DecodeAttentionConfig,
    pub dense_intermediate: usize,
    pub expert_intermediate: usize,
    pub experts: usize,
    pub top_k: usize,
    pub router_norm_multiplier: f32,
    pub activation: GatedActivation,
}

#[derive(Clone, Copy)]
pub struct DecodeMoeBlockWeights<'a> {
    pub attention: DecodeAttentionWeights<'a>,
    pub post_attention_norm: &'a CudaTensor,
    pub pre_dense_norm: &'a CudaTensor,
    pub dense_gate_up: &'a Bf16LinearPairWeights,
    pub dense_down: &'a CudaTensor,
    pub post_dense_norm: &'a CudaTensor,
    pub router: RouterTensors<'a>,
    pub pre_expert_norm: &'a CudaTensor,
    pub post_expert_norm: &'a CudaTensor,
    pub post_feed_forward_norm: &'a CudaTensor,
    pub layer_scalar: &'a CudaTensor,
}

pub(super) fn validate(config: DecodeMoeBlockConfig) -> Result<()> {
    if config.dense_intermediate == 0
        || config.expert_intermediate == 0
        || config.experts == 0
        || config.top_k == 0
        || config.top_k > config.experts
        || !config.router_norm_multiplier.is_finite()
    {
        Err(Error::InvalidDecoderKernel("invalid routed MoE block configuration"))
    } else {
        Ok(())
    }
}