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

use mircuda::{DeviceBuffer, bf16};

use super::GateUpBuffers;
use crate::{
    CudaBackend, CudaTensor, DenseGateUpWeights, Error, MxFp4Bf16Linear, Result, RmsNormBf16,
};

pub(super) fn prepare_gate_up(
    backend: &CudaBackend,
    tokens: usize,
    weights: Option<DenseGateUpWeights<'_>>,
) -> Result<Box<[MxFp4Bf16Linear; 2]>> {
    let DenseGateUpWeights::MxFp4 { gate, up } =
        weights.ok_or(Error::InvalidExecutionPlan("MXFP4 MLP requires gate/up weights"))?
    else {
        return Err(Error::InvalidExecutionPlan("MXFP4 MLP received other gate/up weights"));
    };
    Ok(Box::new([gate.prepare(backend, tokens)?, up.prepare(backend, tokens)?]))
}

pub(super) fn execute_gate_up(
    operations: &[MxFp4Bf16Linear; 2],
    input: &DeviceBuffer<bf16>,
    input_norm: &RmsNormBf16,
    norm_weight: &CudaTensor,
    weights: DenseGateUpWeights<'_>,
    buffers: &mut GateUpBuffers<'_>,
) -> Result<bool> {
    let DenseGateUpWeights::MxFp4 { gate, up } = weights else {
        return Err(Error::InvalidExecutionPlan("MXFP4 gate/up operation received other weights"));
    };
    input_norm.execute(input, norm_weight, buffers.normalized)?;
    for ((operation, weight), output) in
        operations.iter().zip([gate, up]).zip(buffers.separate.iter_mut())
    {
        operation.execute(buffers.normalized, weight, output)?;
    }
    Ok(true)
}