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

mod fixture;

use mircuda::{DeviceBuffer, DeviceElement, bf16};

use super::super::*;
use crate::{
    CudaConfig, ExecutionPhase, PlanSource, Result,
    backend::tuning::{MoeProfileExecution, MoeProfileRequest},
};

#[test]
fn executes_affine_shared_and_routed_experts_for_prefill_and_decode() -> Result<()> {
    let backend = CudaBackend::new(CudaConfig::default())?;
    let config = config();
    let fixture = fixture::MoeFixture::new(config)?;
    let tensors = fixture.upload(&backend)?;
    let moe = CudaAffineSharedExpertMoe::from_tensors(&backend, &tensors, fixture::PREFIX, config)?;

    let input = copy(&backend, &vec![bf16::ZERO; 2 * config.hidden_size])?;
    let mut output = allocate(&backend, 2 * config.hidden_size)?;
    moe.prepare(2)?.execute(&input, &mut output)?;
    assert!(read(&backend, &output)?.iter().all(|value| *value == bf16::ZERO));
    measured(&backend, config, ExecutionPhase::Prefill, 2);

    let input = copy(&backend, &vec![bf16::ZERO; config.hidden_size])?;
    let mut output = allocate(&backend, config.hidden_size)?;
    moe.prepare(1)?.execute(&input, &mut output)?;
    assert!(read(&backend, &output)?.iter().all(|value| *value == bf16::ZERO));
    measured(&backend, config, ExecutionPhase::Decode, 1);
    Ok(())
}

fn measured(
    backend: &CudaBackend,
    config: AffineSharedExpertMoeConfig,
    phase: ExecutionPhase,
    tokens: usize,
) {
    let request = MoeProfileRequest::affine(
        phase,
        tokens,
        config.expert_count,
        config.top_k,
        config.hidden_size,
        config.routed_intermediate_size,
        config.group_size,
        config.expert_bits,
        config.activation,
    );
    assert!(matches!(
        backend.auto_tuner().lookup_moe(request),
        Some((MoeProfileExecution::Affine(_), PlanSource::MeasuredStartup))
    ));
}

fn config() -> AffineSharedExpertMoeConfig {
    AffineSharedExpertMoeConfig {
        hidden_size: 64,
        routed_intermediate_size: 64,
        shared_intermediate_size: 64,
        expert_count: 3,
        top_k: 2,
        group_size: 64,
        expert_bits: 4,
        router_bits: 8,
        activation: GatedActivation::Silu,
    }
}

fn allocate(backend: &CudaBackend, elements: usize) -> Result<DeviceBuffer<bf16>> {
    Ok(backend.inner.pool.allocate(&backend.inner.stream, elements)?)
}

fn copy<T: DeviceElement>(backend: &CudaBackend, values: &[T]) -> Result<DeviceBuffer<T>> {
    let mut host = backend.inner.context.allocate_pinned(values.len())?;
    host.copy_from_slice(values)?;
    let mut device = backend.inner.pool.allocate(&backend.inner.stream, values.len())?;
    backend.inner.stream.copy_to_device(&mut host, &mut device)?;
    Ok(device)
}

fn read<T: DeviceElement>(backend: &CudaBackend, source: &DeviceBuffer<T>) -> Result<Vec<T>> {
    let mut host = backend.inner.context.allocate_pinned(source.len())?;
    backend.inner.stream.copy_to_host(source, &mut host)?;
    Ok(host.to_vec()?)
}