libmir-cuda 0.3.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
86
87

use mircuda::{DeviceBuffer, bf16};
use runtime::kv::{BlockTable, KvWritePlan};

use super::{LayerPrefill, SessionLayer};
use crate::{Error, Result};

impl SessionLayer {
    pub(in crate::backend::model) fn prepare_prefill(&mut self, tokens: usize) -> Result<()> {
        match self {
            Self::Moe(layer) => {
                if !layer.prefill.contains_key(&tokens) {
                    layer.prefill.insert(tokens, layer.template.instantiate_prefill(tokens)?);
                }
            },
            Self::Dense(layer) => {
                if !layer.prefill.contains_key(&tokens) {
                    layer.prefill.insert(tokens, layer.template.instantiate_prefill(tokens)?);
                }
            },
        }
        Ok(())
    }

    pub(in crate::backend::model) fn prefill_plan(
        &mut self,
        tokens: usize,
    ) -> Result<LayerPrefill<'_>> {
        match self {
            Self::Moe(layer) => Ok(LayerPrefill::Moe(
                layer
                    .prefill
                    .get_mut(&tokens)
                    .ok_or(Error::InvalidDecoderKernel("missing CUDA prefill plan"))?,
            )),
            Self::Dense(layer) => Ok(LayerPrefill::Dense(
                layer
                    .prefill
                    .get_mut(&tokens)
                    .ok_or(Error::InvalidDecoderKernel("missing dense CUDA prefill plan"))?,
            )),
        }
    }

    #[allow(clippy::too_many_arguments)]
    pub(in crate::backend::model) fn execute_prefill_masked(
        &mut self,
        input: &DeviceBuffer<bf16>,
        output: &mut DeviceBuffer<bf16>,
        write_plan: &KvWritePlan,
        table: &BlockTable,
        start_position: usize,
        tokens: usize,
        image: Option<crate::backend::attention::ImageAttentionSpan>,
    ) -> Result<()> {
        match self {
            Self::Moe(layer) => {
                let decode = layer.decode.as_mut().ok_or(Error::InvalidDecoderKernel(
                    "CUDA prefill executor belongs to model graph",
                ))?;
                let prefill = layer
                    .prefill
                    .get_mut(&tokens)
                    .ok_or(Error::InvalidDecoderKernel("missing CUDA prefill plan"))?;
                decode.execute_prefill_masked(
                    prefill, input, write_plan, table, start_position, output, image,
                )
            },
            Self::Dense(layer) => {
                if image.is_some() {
                    return Err(Error::UnsupportedVisionContract(
                        "bidirectional pooled-image prefill requires a CUDA hybrid-MoE decoder"
                            .into(),
                    ));
                }
                let weights = layer.template.weights();
                let decode = layer.decode.as_mut().ok_or(Error::InvalidDecoderKernel(
                    "CUDA prefill executor belongs to model graph",
                ))?;
                let prefill = layer
                    .prefill
                    .get_mut(&tokens)
                    .ok_or(Error::InvalidDecoderKernel("missing dense CUDA prefill plan"))?;
                prefill.execute(decode, input, weights, write_plan, table, start_position, output)
            },
        }
    }
}