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

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

pub(super) fn prepare_qkv(
    backend: &CudaBackend,
    tokens: usize,
    weights: Option<DecodeQkvWeights<'_>>,
) -> Result<Box<[MxFp8Bf16Linear; 3]>> {
    let DecodeQkvWeights::MxFp8(weights) =
        weights.ok_or(Error::InvalidExecutionPlan("MXFP8 attention requires prepared QKV"))?
    else {
        return Err(Error::InvalidExecutionPlan("MXFP8 attention received other QKV weights"));
    };
    Ok(Box::new([
        weights[0].prepare(backend, tokens)?,
        weights[1].prepare(backend, tokens)?,
        weights[2].prepare(backend, tokens)?,
    ]))
}