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
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

use super::SharedRoutedExecution;
use crate::{
    CudaBackend, DenseRole, Result,
    backend::{
        linear::{CheckpointProjection, MxFp4ExpertWeights, MxFp4GatheredMoeBf16},
        shared_moe::{AffineSharedExpertMoeConfig, weights::AffineSharedExpertMoeWeights},
    },
    kernels::{RouterUnitSpec, RouterUnitTopK},
};

pub(super) fn mxfp4(
    backend: &CudaBackend,
    config: AffineSharedExpertMoeConfig,
    weights: &AffineSharedExpertMoeWeights,
    experts: &MxFp4ExpertWeights,
    tokens: usize,
) -> Result<SharedRoutedExecution> {
    let (router, top_k, scores, selected, routing) = routing(backend, config, weights, tokens, 4)?;
    Ok(SharedRoutedExecution::MxFp4 {
        router,
        top_k,
        scores,
        selected,
        routing,
        experts: Box::new(MxFp4GatheredMoeBf16::new(
            backend,
            tokens,
            config.top_k,
            config.activation,
            experts,
        )?),
    })
}

pub(super) fn mxfp8(
    backend: &CudaBackend,
    config: AffineSharedExpertMoeConfig,
    weights: &AffineSharedExpertMoeWeights,
    experts: &crate::backend::linear::MxFp8ExpertWeights,
    tokens: usize,
) -> Result<SharedRoutedExecution> {
    let (router, top_k, scores, selected, routing) = routing(backend, config, weights, tokens, 8)?;
    Ok(SharedRoutedExecution::MxFp8 {
        router,
        top_k,
        scores,
        selected,
        routing,
        experts: Box::new(crate::backend::linear::MxFp8GatheredMoeBf16::new(
            backend,
            tokens,
            config.top_k,
            config.activation,
            experts,
        )?),
    })
}

type Routing = (
    Box<CheckpointProjection>,
    RouterUnitTopK,
    mircuda::DeviceBuffer<mircuda::bf16>,
    mircuda::DeviceBuffer<u32>,
    mircuda::DeviceBuffer<mircuda::bf16>,
);

pub(super) fn routing(
    backend: &CudaBackend,
    config: AffineSharedExpertMoeConfig,
    weights: &AffineSharedExpertMoeWeights,
    tokens: usize,
    bits: usize,
) -> Result<Routing> {
    let selections = tokens.checked_mul(config.top_k).ok_or(crate::Error::InvalidDecoderKernel(
        if bits == 4 {
            "MXFP4 shared routing size overflow"
        } else {
            "MXFP8 shared routing size overflow"
        },
    ))?;
    Ok((
        Box::new(CheckpointProjection::new(
            backend,
            tokens,
            config.hidden_size,
            config.expert_count,
            DenseRole::Router,
            &weights.router,
        )?),
        RouterUnitTopK::compile(
            backend.compiler(),
            RouterUnitSpec {
                tokens,
                experts: config.expert_count,
                top_k: config.top_k,
            },
        )?,
        backend.pool().allocate(backend.stream(), tokens * config.expert_count)?,
        backend.pool().allocate(backend.stream(), selections)?,
        backend.pool().allocate(backend.stream(), selections)?,
    ))
}