turbospark-runtime 0.3.0

The raw-completion prefill+decode loop: wires a LogitProducer, the tokenizer, and the selection sampler into one token generation loop.
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
103
104
105
106
107
108
109
110
111
112
113
114

//! Family-specific router arithmetic shared by decode and chunked prefill.
use super::{layer_tensor, state::RealLlamaState};
use crate::real_forward::RealForwardError;
use crate::real_forward_utils::entry;
use model_io::ResidentIndex;

pub(crate) fn fp32_tensor(name: &str) -> bool {
    name.starts_with("language_model.model.layers.")
        && (name.ends_with(".mlp.gate.weight") || name.ends_with(".mlp.e_score_correction_bias"))
}

#[allow(clippy::too_many_arguments)]
pub(super) fn encode(
    context: &mut gpu::MetalContext,
    pass: &gpu::PassEncoder,
    weights: &gpu::ResidentGpuWeights,
    index: &ResidentIndex,
    state: &RealLlamaState,
    layer: usize,
    token: usize,
    hidden: usize,
    experts: usize,
) -> Result<(), RealForwardError> {
    let name = layer_tensor(layer, "mlp.gate.weight");
    let router = entry(index, &name)?;
    let fp32 = !state.correction_bias.is_empty();
    let dtype = if fp32 { 3 } else { 5 };
    let size = experts * hidden * if fp32 { 4 } else { 1 };
    if router.dtype != dtype || router.size_bytes as usize != size {
        return Err(RealForwardError::Unsupported(format!(
            "{name}: expected dtype {dtype} and {size} bytes"
        )));
    }
    let base = index.header.index_size;
    let view = |off| (weights.buffer(), weights.gpu_offset(off - base));
    let x = (&state.moe_x, (token * hidden * 2) as u64);
    let out = (&state.router_logits_f32, (token * experts * 4) as u64);
    if fp32 {
        gpu::encode_minimax_router(
            context,
            pass,
            view(router.file_offset),
            x,
            out,
            experts as u32,
            hidden as u32,
        )
        .map_err(RealForwardError::Gpu)?;
    } else {
        gpu::encode_router_gemv_gemma4(
            context,
            pass,
            view(router.file_offset),
            view(router.scale_offset),
            view(router.bias_offset),
            x,
            (&state.router_ones, 0),
            out,
            experts as u32,
            hidden as u32,
        )
        .map_err(RealForwardError::Gpu)?;
    }
    Ok(())
}

pub(crate) fn sigmoid_topk(logits: &[f32], bias: &[f32], k: usize) -> (Vec<usize>, Vec<f32>) {
    assert_eq!(logits.len(), bias.len());
    let scores: Vec<f32> = logits
        .iter()
        .map(|&x| {
            if x >= 0.0 {
                1.0 / (1.0 + (-x).exp())
            } else {
                let e = x.exp();
                e / (1.0 + e)
            }
        })
        .collect();
    let mut order: Vec<usize> = (0..scores.len()).collect();
    order.sort_by(|&a, &b| {
        (scores[b] + bias[b])
            .total_cmp(&(scores[a] + bias[a]))
            .then(a.cmp(&b))
    });
    order.truncate(k);
    let sum = order.iter().map(|&i| scores[i]).sum::<f32>() + 1e-20;
    let weights = order.iter().map(|&i| scores[i] / sum).collect();
    (order, weights)
}

#[cfg(test)]
mod tests {
    use super::sigmoid_topk;
    #[test]
    fn correction_selects_but_does_not_weight() {
        let (ids, w) = sigmoid_topk(&[0.0, 3f32.ln(), -3f32.ln()], &[0.0, 0.0, 1.0], 2);
        assert_eq!(ids, [2, 1]);
        assert!((w[0] - 0.25).abs() < 1e-7);
        assert!((w[1] - 0.75).abs() < 1e-7);
        // Exercise the correction on both sides of the sort comparator.
        let (ids, w) = sigmoid_topk(&[-3f32.ln(), 3f32.ln(), 0.0], &[1.0, 0.0, 0.0], 2);
        assert_eq!(ids, [0, 1]);
        assert!((w[0] - 0.25).abs() < 1e-7);
    }
    #[test]
    fn sigmoid_weights_are_not_softmax_and_ties_are_stable() {
        let (ids, w) = sigmoid_topk(&[0.0, 0.0, 3f32.ln()], &[0.0; 3], 3);
        assert_eq!(ids, [2, 0, 1]);
        assert!((w[0] - 3.0 / 7.0).abs() < 1e-7);
        assert!((w[2] - 2.0 / 7.0).abs() < 1e-7);
        assert_eq!(sigmoid_topk(&[0.0; 3], &[0.0; 3], 2).0, [0, 1]);
    }
}