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

use mircuda::{DeviceBuffer, bf16};

use crate::{CudaBackend, Error, Result};

#[derive(Debug)]
pub(super) struct FullAttentionLayerScratch {
    pub(super) normalized: DeviceBuffer<bf16>,
    pub(super) attention: DeviceBuffer<bf16>,
    pub(super) residual: DeviceBuffer<bf16>,
    pub(super) moe: DeviceBuffer<bf16>,
}

impl FullAttentionLayerScratch {
    pub(super) fn new(backend: &CudaBackend, tokens: usize, hidden: usize) -> Result<Self> {
        let elements = tokens
            .checked_mul(hidden)
            .ok_or(Error::InvalidDecoderKernel("gated full-attention layer scratch overflow"))?;
        let allocate = || backend.inner.pool.allocate(&backend.inner.stream, elements);
        Ok(Self {
            normalized: allocate()?,
            attention: allocate()?,
            residual: allocate()?,
            moe: allocate()?,
        })
    }
}