singe_kernel/cpu/
moe.rs

1//! Mixture-of-experts routing and alignment helpers.
2
3pub fn moe_align_block_size_i32(
4    topk_ids: &[i32],
5    expert_count: usize,
6    block_size: usize,
7) -> (Vec<i32>, Vec<i32>, i32, Vec<i32>, i32) {
8    let sentinel = topk_ids.len() as i32;
9    let max_num_tokens_padded = topk_ids.len() + expert_count * (block_size - 1);
10    let mut sorted_token_ids = vec![sentinel; max_num_tokens_padded];
11    let mut cumsum = vec![0i32; expert_count + 1];
12    let mut total = 0usize;
13    let mut max_count = 0usize;
14    for (expert, cumsum_slot) in cumsum.iter_mut().enumerate().take(expert_count) {
15        *cumsum_slot = total as i32;
16        let mut count = 0usize;
17        for (token, &id) in topk_ids.iter().enumerate() {
18            if id == expert as i32 {
19                sorted_token_ids[total + count] = token as i32;
20                count += 1;
21            }
22        }
23        max_count = max_count.max(count);
24        total += ceil_div(count, block_size) * block_size;
25    }
26    cumsum[expert_count] = total as i32;
27    let expert_ids = (0..total / block_size)
28        .map(|block| {
29            let token_offset = (block * block_size) as i32;
30            cumsum
31                .windows(2)
32                .position(|range| token_offset >= range[0] && token_offset < range[1])
33                .unwrap_or(expert_count - 1) as i32
34        })
35        .collect::<Vec<_>>();
36    (
37        sorted_token_ids,
38        expert_ids,
39        total as i32,
40        cumsum,
41        max_count as i32,
42    )
43}
44
45pub fn fused_moe_f32(
46    input: &[f32],
47    weight: &[f32],
48    routed_weight: &[f32],
49    topk_ids: &[i32],
50    tokens: usize,
51    top_k: usize,
52    columns: usize,
53    reduction: usize,
54    mul_routed_weight: bool,
55) -> Vec<f32> {
56    let mut out = vec![0.0f32; tokens * top_k * columns];
57    for slot in 0..tokens * top_k {
58        let token = slot / top_k;
59        let expert = topk_ids[slot] as usize;
60        for column in 0..columns {
61            let mut sum = 0.0f32;
62            for reduction_index in 0..reduction {
63                let input_offset = token * reduction + reduction_index;
64                let weight_offset =
65                    expert * columns * reduction + column * reduction + reduction_index;
66                sum += input[input_offset] * weight[weight_offset];
67            }
68            if mul_routed_weight {
69                sum *= routed_weight[slot];
70            }
71            out[slot * columns + column] = sum;
72        }
73    }
74    out
75}
76
77pub fn fused_moe_f8e4m3_block_scaled_f32(
78    input: &[u8],
79    weight: &[u8],
80    input_scales: &[f32],
81    weight_scales: &[f32],
82    routed_weight: &[f32],
83    topk_ids: &[i32],
84    tokens: usize,
85    top_k: usize,
86    _experts: usize,
87    columns: usize,
88    reduction: usize,
89    group_n: usize,
90    group_k: usize,
91    mul_routed_weight: bool,
92) -> Vec<f32> {
93    let k_groups = ceil_div(reduction, group_k);
94    let n_groups = ceil_div(columns, group_n);
95    let mut out = vec![0.0f32; tokens * top_k * columns];
96    for slot in 0..tokens * top_k {
97        let token = slot / top_k;
98        let expert = topk_ids[slot] as usize;
99        for column in 0..columns {
100            let mut sum = 0.0f32;
101            for reduction_index in 0..reduction {
102                let k_group = reduction_index / group_k;
103                let input_offset = token * reduction + reduction_index;
104                let weight_offset =
105                    expert * columns * reduction + column * reduction + reduction_index;
106                let input_scale = input_scales[token * k_groups + k_group];
107                let weight_scale = weight_scales
108                    [expert * n_groups * k_groups + (column / group_n) * k_groups + k_group];
109                sum += f8e4m3_value(input[input_offset])
110                    * f8e4m3_value(weight[weight_offset])
111                    * input_scale
112                    * weight_scale;
113            }
114            if mul_routed_weight {
115                sum *= routed_weight[slot];
116            }
117            out[slot * columns + column] = sum;
118        }
119    }
120    out
121}
122
123pub fn ceil_div(lhs: usize, rhs: usize) -> usize {
124    lhs.div_ceil(rhs)
125}
126
127fn f8e4m3_value(value: u8) -> f32 {
128    let sign = if value & 0x80 == 0 { 1.0 } else { -1.0 };
129    let exponent = (value >> 3) & 0x0f;
130    let mantissa = value & 0x07;
131    if exponent == 0x0f && mantissa == 0x07 {
132        f32::NAN
133    } else if exponent == 0 {
134        if mantissa == 0 {
135            sign * 0.0
136        } else {
137            sign * (mantissa as f32 / 8.0) * 2f32.powi(-6)
138        }
139    } else {
140        sign * (1.0 + mantissa as f32 / 8.0) * 2f32.powi(exponent as i32 - 7)
141    }
142}
singe_kernel/cpu/moe.rs

singe_kernel/cpu/
moe.rs
