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

use mircuda::{DeviceBuffer, bf16};

use super::GroupedNvFp4LinearBf16;
use crate::{CudaBackend, NvFp4ExpertBank, Result};

#[derive(Debug)]
pub(in crate::backend::linear::nvfp4::selected) struct GroupedNvFp4PairBf16 {
    left: GroupedNvFp4LinearBf16,
    right: GroupedNvFp4LinearBf16,
}

impl GroupedNvFp4PairBf16 {
    pub(in crate::backend::linear::nvfp4::selected) fn new(
        backend: &CudaBackend,
        tokens: usize,
        selected: usize,
        left: NvFp4ExpertBank,
        right: NvFp4ExpertBank,
    ) -> Result<Self> {
        Ok(Self {
            left: GroupedNvFp4LinearBf16::new(backend, tokens, selected, left)?,
            right: GroupedNvFp4LinearBf16::new(backend, tokens, selected, right)?,
        })
    }

    pub(in crate::backend::linear::nvfp4::selected) fn execute(
        &mut self,
        input: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        left_output: &mut DeviceBuffer<bf16>,
        right_output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        GroupedNvFp4LinearBf16::quantize_pair(&mut self.left, &mut self.right, input, selected)?;
        self.left.execute_prepared(selected, left_output)?;
        self.right.execute_prepared(selected, right_output)
    }

    pub(in crate::backend::linear::nvfp4::selected) fn output_elements(&self) -> Result<usize> {
        self.left.output_elements()
    }

    pub(in crate::backend::linear::nvfp4::selected) const fn output_features(&self) -> usize {
        self.left.output_features()
    }
}