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
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

use super::DecodeMoeBlockWeights;
use crate::{
    Bf16LinearPairWeights, CudaTensor, RouterTensors,
    backend::attention::graph::CapturedAttentionWeights,
};

#[derive(Clone)]
pub(in crate::backend) struct CapturedBlockWeights {
    attention: CapturedAttentionWeights,
    post_attention_norm: CudaTensor,
    pre_dense_norm: CudaTensor,
    dense_gate_up: Bf16LinearPairWeights,
    dense_down: CudaTensor,
    post_dense_norm: CudaTensor,
    router_projection: CudaTensor,
    router_norm_scale: CudaTensor,
    router_expert_scale: CudaTensor,
    pre_expert_norm: CudaTensor,
    post_expert_norm: CudaTensor,
    post_feed_forward_norm: CudaTensor,
    layer_scalar: CudaTensor,
}

impl CapturedBlockWeights {
    pub(in crate::backend) fn prepare(
        backend: &crate::CudaBackend,
        weights: DecodeMoeBlockWeights<'_>,
    ) -> crate::Result<Self> {
        Ok(Self {
            attention: CapturedAttentionWeights::prepare(backend, weights.attention)?,
            post_attention_norm: weights.post_attention_norm.clone(),
            pre_dense_norm: weights.pre_dense_norm.clone(),
            dense_gate_up: weights.dense_gate_up.clone(),
            dense_down: weights.dense_down.clone(),
            post_dense_norm: weights.post_dense_norm.clone(),
            router_projection: weights.router.projection.clone(),
            router_norm_scale: weights.router.norm_scale.clone(),
            router_expert_scale: weights.router.expert_scale.clone(),
            pre_expert_norm: weights.pre_expert_norm.clone(),
            post_expert_norm: weights.post_expert_norm.clone(),
            post_feed_forward_norm: weights.post_feed_forward_norm.clone(),
            layer_scalar: weights.layer_scalar.clone(),
        })
    }

    pub(in crate::backend::block) fn borrow(&self) -> DecodeMoeBlockWeights<'_> {
        DecodeMoeBlockWeights {
            attention: self.attention.borrow(),
            post_attention_norm: &self.post_attention_norm,
            pre_dense_norm: &self.pre_dense_norm,
            dense_gate_up: &self.dense_gate_up,
            dense_down: &self.dense_down,
            post_dense_norm: &self.post_dense_norm,
            router: RouterTensors {
                projection: &self.router_projection,
                norm_scale: &self.router_norm_scale,
                expert_scale: &self.router_expert_scale,
            },
            pre_expert_norm: &self.pre_expert_norm,
            post_expert_norm: &self.post_expert_norm,
            post_feed_forward_norm: &self.post_feed_forward_norm,
            layer_scalar: &self.layer_scalar,
        }
    }
}

impl From<DecodeMoeBlockWeights<'_>> for CapturedBlockWeights {
    fn from(weights: DecodeMoeBlockWeights<'_>) -> Self {
        Self {
            attention: weights.attention.into(),
            post_attention_norm: weights.post_attention_norm.clone(),
            pre_dense_norm: weights.pre_dense_norm.clone(),
            dense_gate_up: weights.dense_gate_up.clone(),
            dense_down: weights.dense_down.clone(),
            post_dense_norm: weights.post_dense_norm.clone(),
            router_projection: weights.router.projection.clone(),
            router_norm_scale: weights.router.norm_scale.clone(),
            router_expert_scale: weights.router.expert_scale.clone(),
            pre_expert_norm: weights.pre_expert_norm.clone(),
            post_expert_norm: weights.post_expert_norm.clone(),
            post_feed_forward_norm: weights.post_feed_forward_norm.clone(),
            layer_scalar: weights.layer_scalar.clone(),
        }
    }
}