ruDNN 0.1.57

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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

mod dispatch;
mod experts;
mod kernels;
mod routing;
mod received;
mod projection;
mod activation;
pub use projection::{ExpertProjectionCache,expert_projection};
pub use activation::{SwiGluActivationSelection,SwiGluActivationBackward,swiglu_activation,swiglu_activation_backward};
#[cfg(feature="tensor-nf4-moe")]
mod nf4;
#[cfg(feature="tensor-nf4-moe")]
pub use nf4::{Nf4ExpertProjection,Nf4ExpertError,Nf4ExpertExecution,Nf4SwiGluExperts,Nf4SwiGluCache};
#[cfg(feature="tensor-packed-moe")]
mod packed;
#[cfg(feature="tensor-packed-moe")]
pub use packed::{PackedExpertError,PackedExpertProjection,PackedSwiGluExperts,PackedSwiGluCache};

pub use dispatch::{DispatchedTokens,CombineBackward,CombineGradientStrategy,CombineGradientSelection,CombineBackwardSelected};
pub use experts::{SwiGluExperts,ExpertTrainingCache,ExpertTrainingOutput,ExpertBackward,ExpertGradientSelection,ExpertBackwardSelected,GroupedExpertRows};
pub use received::ReceivedExpertRows;
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};

mod router_training;
pub use router_training::{RouterScoring, RouterWeightOptions, RouterTrainingPlan,
    selected_router_weights, selected_router_weights_into,
    selected_router_backward, selected_router_backward_into};

#[cfg(test)]
mod tests_v31;