libmir-cuda 0.1.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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

use super::{
    projection::OutputHeadWeight,
    weight::{CudaOutputHeadTemplate, bf16},
};
use crate::{
    CudaBackend, CudaTensor, OutputHeadExecution, Result,
    kernels::{BlockFp8LinearKernels, BlockFp8LinearSpec, Fp8OutputSpec, Fp8RefinementKernels},
};

pub(super) fn prepare(
    backend: &CudaBackend,
    source: &CudaTensor,
    input_features: usize,
    output_features: usize,
    execution: OutputHeadExecution,
) -> Result<CudaOutputHeadTemplate> {
    let source_weight = bf16(source)?;
    let spec = BlockFp8LinearSpec::new(input_features, output_features)?;
    let kernels = BlockFp8LinearKernels::compile(&backend.inner.compiler, spec)?;
    let mut weight = backend
        .inner
        .pool
        .allocate::<u8>(&backend.inner.stream, spec.weight_elements()?)?;
    let mut scales = backend
        .inner
        .pool
        .allocate::<f32>(&backend.inner.stream, spec.scale_elements()?)?;
    kernels.quantize(&backend.inner.stream, source_weight, &mut weight, &mut scales)?;
    tracing::debug!(
        input_features,
        output_features,
        scale_block = 128,
        "prepared block-scaled E4M3 CUDA output head"
    );
    let weight = match execution {
        OutputHeadExecution::Fp8BlockVectorized => {
            OutputHeadWeight::Fp8BlockVectorized { kernels, weight, scales }
        },
        OutputHeadExecution::AutoRefined | OutputHeadExecution::Fp8BlockRefined => {
            tracing::debug!(
                candidates = Fp8RefinementKernels::candidate_count(),
                "prepared exact BF16 output candidate refinement"
            );
            OutputHeadWeight::Fp8BlockRefined {
                refinement: Fp8RefinementKernels::compile(
                    &backend.inner.compiler,
                    Fp8OutputSpec::new(input_features, output_features)?,
                )?,
                exact_tensor: source.clone(),
                exact_weight: source_weight.clone(),
                kernels,
                weight,
                scales,
            }
        },
        _ => {
            return Err(crate::Error::InvalidExecutionPlan(
                "non-block output execution reached block preparation",
            ));
        },
    };
    Ok(CudaOutputHeadTemplate {
        exact: source.clone(),
        weight,
        kernels: None,
        execution,
        input_features,
        output_features,
    })
}