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

use mircuda::{DeviceBuffer, bf16};

use super::weight::CudaOutputHeadTemplate;
use crate::{
    Bf16Projection, CudaBackend, CudaTensor, DensePlanRequest, DenseRole, ExecutionPhase, Result,
};

/// Exact BF16 output projection prepared for one fixed decode bucket.
#[derive(Debug)]
pub struct CudaBatchOutputHead {
    projection: Bf16Projection,
    weight: CudaTensor,
}

impl CudaBatchOutputHead {
    pub(crate) fn new(
        backend: &CudaBackend,
        template: &CudaOutputHeadTemplate,
        rows: usize,
    ) -> Result<Self> {
        Ok(Self {
            projection: backend.prepare_bf16_projection(DensePlanRequest {
                phase: ExecutionPhase::Decode,
                role: DenseRole::OutputHead,
                tokens: rows,
                input_features: template.input_features,
                output_features: template.output_features,
            })?,
            weight: template.exact.clone(),
        })
    }

    pub(crate) fn execute(
        &mut self,
        input: &DeviceBuffer<bf16>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        self.projection.execute(input, &self.weight, output)
    }
}