ruDNN 0.1.7

Neural network operators for Ruda.
Documentation 
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

mod dispatch;
mod experts;
mod kernels;
mod routing;

pub use dispatch::DispatchedTokens;
pub use experts::SwiGluExperts;
pub use routing::{RoutingOptions, RoutingPlan, route};

use ruda_core::{
    device::Device,
    tensor::{DType, Shape},
};
use ruda_kernel::{
    dsl::Runtime,
    tensor::{RudaTensor, allocation::empty_device_contiguous_dtype},
};
use std::fmt;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MoeError(pub &'static str);

impl fmt::Display for MoeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.0)
    }
}

impl std::error::Error for MoeError {}

impl From<rublas::tensor_grouped::GroupedMatmulError> for MoeError {
    fn from(error: rublas::tensor_grouped::GroupedMatmulError) -> Self {
        Self(error.0)
    }
}

fn elements(dimensions: &[usize]) -> Result<usize, MoeError> {
    if dimensions.iter().any(|&dim| dim > u32::MAX as usize) {
        return Err(MoeError("MoE dimension exceeds U32 indexing"));
    }
    dimensions
        .iter()
        .try_fold(1usize, |size, &dim| size.checked_mul(dim))
        .filter(|&size| size <= u32::MAX as usize)
        .ok_or(MoeError("MoE tensor exceeds U32 indexing"))
}

fn same_device<R: Runtime>(a: &RudaTensor<R>, b: &RudaTensor<R>) -> Result<(), MoeError> {
    if a.device.to_id() != b.device.to_id() {
        return Err(MoeError("MoE operands must share a device"));
    }
    Ok(())
}

fn float_tensor<R: Runtime>(tensor: &RudaTensor<R>) -> Result<(), MoeError> {
    if tensor.qparams.is_some() || !matches!(tensor.dtype, DType::F32 | DType::F16 | DType::BF16) {
        return Err(MoeError(
            "MoE requires non-quantized F32, F16 or BF16 tensors",
        ));
    }
    Ok(())
}

fn empty<R: Runtime>(like: &RudaTensor<R>, shape: impl Into<Shape>, dtype: DType) -> RudaTensor<R> {
    empty_device_contiguous_dtype(
        like.client.clone(),
        like.device.clone(),
        shape.into(),
        dtype,
    )
}

#[cfg(test)]
mod tests;

/// Public expert-matrix execution policy.
pub use rublas::tensor_grouped::GroupedStrategy;

mod grouped_routing;
pub use grouped_routing::{GroupRoutingOptions,route_sigmoid_grouped};