skippy_protocol/binary/
activation.rs

1use std::io;
2
3use super::{
4    WireActivationDType, invalid_data, state_flags, types::MAX_STAGE_DECODED_ACTIVATION_BYTES,
5};
6
7pub fn activation_wire_bytes(
8    dtype: WireActivationDType,
9    token_count: i32,
10    n_embd: i32,
11) -> io::Result<usize> {
12    activation_wire_bytes_with_state_flags(dtype, token_count, n_embd, 0)
13}
14
15pub fn activation_wire_bytes_with_state_flags(
16    dtype: WireActivationDType,
17    token_count: i32,
18    n_embd: i32,
19    state_flag_bits: i32,
20) -> io::Result<usize> {
21    if token_count < 0 || n_embd < 0 {
22        return Err(invalid_data("negative activation dimensions"));
23    }
24    let token_count = (token_count as usize)
25        .checked_mul(activation_payload_multiplier_from_state_flags(
26            state_flag_bits,
27        ))
28        .ok_or_else(|| invalid_data("activation token count overflow"))?;
29    let n_embd = n_embd as usize;
30    let elements = token_count
31        .checked_mul(n_embd)
32        .ok_or_else(|| invalid_data("activation element count overflow"))?;
33    match dtype {
34        WireActivationDType::F32 => elements
35            .checked_mul(4)
36            .ok_or_else(|| invalid_data("activation byte count overflow")),
37        WireActivationDType::F16 => elements
38            .checked_mul(2)
39            .ok_or_else(|| invalid_data("activation byte count overflow")),
40        WireActivationDType::Q8 => token_count
41            .checked_mul(4)
42            .and_then(|scales| scales.checked_add(elements))
43            .ok_or_else(|| invalid_data("activation byte count overflow")),
44    }
45}
46
47pub(crate) fn activation_decoded_f32_bytes_with_state_flags(
48    token_count: i32,
49    n_embd: i32,
50    state_flag_bits: i32,
51) -> io::Result<usize> {
52    activation_wire_bytes_with_state_flags(
53        WireActivationDType::F32,
54        token_count,
55        n_embd,
56        state_flag_bits,
57    )
58}
59
60pub fn encode_f32_activation_payload(
61    dtype: WireActivationDType,
62    token_count: i32,
63    n_embd: i32,
64    f32_payload: &[u8],
65) -> io::Result<Vec<u8>> {
66    encode_f32_activation_payload_with_state_flags(dtype, token_count, n_embd, f32_payload, 0)
67}
68
69pub fn encode_f32_activation_payload_with_state_flags(
70    dtype: WireActivationDType,
71    token_count: i32,
72    n_embd: i32,
73    f32_payload: &[u8],
74    state_flag_bits: i32,
75) -> io::Result<Vec<u8>> {
76    let expected_f32_bytes = activation_wire_bytes_with_state_flags(
77        WireActivationDType::F32,
78        token_count,
79        n_embd,
80        state_flag_bits,
81    )?;
82    if expected_f32_bytes > MAX_STAGE_DECODED_ACTIVATION_BYTES {
83        return Err(invalid_data(
84            "decoded activation payload byte count exceeds maximum",
85        ));
86    }
87    if f32_payload.len() != expected_f32_bytes {
88        return Err(invalid_data("F32 activation payload size mismatch"));
89    }
90    match dtype {
91        WireActivationDType::F32 => Ok(f32_payload.to_vec()),
92        WireActivationDType::F16 => encode_f32_to_f16_bytes(f32_payload),
93        WireActivationDType::Q8 => {
94            let token_count = token_count
95                .checked_mul(activation_payload_multiplier_from_state_flags(state_flag_bits) as i32)
96                .ok_or_else(|| invalid_data("Q8 activation token count overflow"))?;
97            encode_f32_to_q8_bytes(f32_payload, token_count, n_embd)
98        }
99    }
100}
101
102pub fn activation_payload_multiplier_from_state_flags(state_flag_bits: i32) -> usize {
103    if (state_flag_bits & state_flags::GEMMA3N_ALTUP_SIDEBAND) != 0 {
104        4
105    } else if (state_flag_bits
106        & (state_flags::INKLING_MTP_EMBD_SIDEBAND | state_flags::RWKV7_V_FIRST_SIDEBAND))
107        != 0
108    {
109        2
110    } else {
111        1
112    }
113}
114
115pub(crate) fn decode_f16_to_f32_bytes(input: &[u8]) -> io::Result<Vec<u8>> {
116    if input.len() & 1 != 0 {
117        return Err(invalid_data("F16 activation payload has odd byte length"));
118    }
119    let decoded_bytes = input
120        .len()
121        .checked_mul(2)
122        .ok_or_else(|| invalid_data("decoded activation byte count overflow"))?;
123    if decoded_bytes > MAX_STAGE_DECODED_ACTIVATION_BYTES {
124        return Err(invalid_data(
125            "decoded activation payload byte count exceeds maximum",
126        ));
127    }
128    let mut out = Vec::with_capacity(decoded_bytes);
129    for chunk in input.chunks_exact(2) {
130        let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
131        out.extend_from_slice(&f16_bits_to_f32(bits).to_le_bytes());
132    }
133    Ok(out)
134}
135
136fn encode_f32_to_f16_bytes(input: &[u8]) -> io::Result<Vec<u8>> {
137    if input.len() & 3 != 0 {
138        return Err(invalid_data("F32 activation payload size is not aligned"));
139    }
140    let mut out = Vec::with_capacity(input.len() / 2);
141    for chunk in input.chunks_exact(4) {
142        let value = f32::from_le_bytes(chunk.try_into().expect("chunks_exact size"));
143        out.extend_from_slice(&f32_to_f16_bits(value).to_le_bytes());
144    }
145    Ok(out)
146}
147
148pub(crate) fn decode_q8_to_f32_bytes_with_state_flags(
149    input: &[u8],
150    token_count: i32,
151    n_embd: i32,
152    state_flag_bits: i32,
153) -> io::Result<Vec<u8>> {
154    if token_count < 0 || n_embd < 0 {
155        return Err(invalid_data("negative Q8 activation dimensions"));
156    }
157    let token_count = (token_count as usize)
158        .checked_mul(activation_payload_multiplier_from_state_flags(
159            state_flag_bits,
160        ))
161        .ok_or_else(|| invalid_data("Q8 activation token count overflow"))?;
162    let n_embd = n_embd as usize;
163    let scale_bytes = token_count
164        .checked_mul(4)
165        .ok_or_else(|| invalid_data("Q8 scale byte count overflow"))?;
166    let value_bytes = token_count
167        .checked_mul(n_embd)
168        .ok_or_else(|| invalid_data("Q8 value byte count overflow"))?;
169    let expected_bytes = scale_bytes
170        .checked_add(value_bytes)
171        .ok_or_else(|| invalid_data("Q8 activation payload byte count overflow"))?;
172    if input.len() != expected_bytes {
173        return Err(invalid_data("Q8 activation payload size mismatch"));
174    }
175    let decoded_bytes = value_bytes
176        .checked_mul(4)
177        .ok_or_else(|| invalid_data("decoded activation byte count overflow"))?;
178    if decoded_bytes > MAX_STAGE_DECODED_ACTIVATION_BYTES {
179        return Err(invalid_data(
180            "decoded activation payload byte count exceeds maximum",
181        ));
182    }
183    let mut out = Vec::with_capacity(decoded_bytes);
184    for token_index in 0..token_count {
185        let scale_offset = token_index * 4;
186        let scale = f32::from_le_bytes([
187            input[scale_offset],
188            input[scale_offset + 1],
189            input[scale_offset + 2],
190            input[scale_offset + 3],
191        ]);
192        let row_offset = scale_bytes + token_index * n_embd;
193        for value in &input[row_offset..row_offset + n_embd] {
194            let signed = *value as i8;
195            out.extend_from_slice(&((signed as f32) * scale).to_le_bytes());
196        }
197    }
198    Ok(out)
199}
200
201#[cfg(test)]
202pub(crate) fn decode_q8_to_f32_bytes(
203    input: &[u8],
204    token_count: i32,
205    n_embd: i32,
206) -> io::Result<Vec<u8>> {
207    decode_q8_to_f32_bytes_with_state_flags(input, token_count, n_embd, 0)
208}
209
210fn encode_f32_to_q8_bytes(input: &[u8], token_count: i32, n_embd: i32) -> io::Result<Vec<u8>> {
211    if token_count < 0 || n_embd < 0 {
212        return Err(invalid_data("negative Q8 activation dimensions"));
213    }
214    let token_count = token_count as usize;
215    let n_embd = n_embd as usize;
216    let expected_bytes = token_count
217        .checked_mul(n_embd)
218        .and_then(|elements| elements.checked_mul(4))
219        .ok_or_else(|| invalid_data("Q8 source byte count overflow"))?;
220    if input.len() != expected_bytes {
221        return Err(invalid_data("Q8 source payload size mismatch"));
222    }
223
224    let mut scales = Vec::with_capacity(token_count * 4);
225    let mut packed = Vec::with_capacity(token_count * n_embd);
226    for token_index in 0..token_count {
227        let row_offset = token_index * n_embd * 4;
228        let row = &input[row_offset..row_offset + n_embd * 4];
229        let mut max_abs = 0.0_f32;
230        for chunk in row.chunks_exact(4) {
231            let value = f32::from_le_bytes(chunk.try_into().expect("chunks_exact size"));
232            max_abs = max_abs.max(value.abs());
233        }
234        let scale = if max_abs > 0.0 { max_abs / 127.0 } else { 1.0 };
235        scales.extend_from_slice(&scale.to_le_bytes());
236        for chunk in row.chunks_exact(4) {
237            let value = f32::from_le_bytes(chunk.try_into().expect("chunks_exact size"));
238            let quantized = (value / scale).round().clamp(-127.0, 127.0) as i8;
239            packed.push(quantized as u8);
240        }
241    }
242    scales.extend_from_slice(&packed);
243    Ok(scales)
244}
245
246fn f16_bits_to_f32(bits: u16) -> f32 {
247    let sign = ((bits & 0x8000) as u32) << 16;
248    let exponent = (bits >> 10) & 0x1f;
249    let mantissa = bits & 0x03ff;
250    let f32_bits = if exponent == 0 {
251        if mantissa == 0 {
252            sign
253        } else {
254            let mut mant = mantissa as u32;
255            let mut exp = -14_i32;
256            while (mant & 0x0400) == 0 {
257                mant <<= 1;
258                exp -= 1;
259            }
260            mant &= 0x03ff;
261            let exp_bits = ((exp + 127) as u32) << 23;
262            sign | exp_bits | (mant << 13)
263        }
264    } else if exponent == 0x1f {
265        sign | 0x7f80_0000 | ((mantissa as u32) << 13)
266    } else {
267        let exp_bits = ((exponent as u32) + (127 - 15)) << 23;
268        sign | exp_bits | ((mantissa as u32) << 13)
269    };
270    f32::from_bits(f32_bits)
271}
272
273fn f32_to_f16_bits(value: f32) -> u16 {
274    let bits = value.to_bits();
275    let sign = ((bits >> 16) & 0x8000) as u16;
276    let exponent = ((bits >> 23) & 0xff) as i32;
277    let mantissa = bits & 0x007f_ffff;
278
279    if exponent == 0xff {
280        let nan_bit = if mantissa == 0 { 0 } else { 0x0200 };
281        return sign | 0x7c00 | nan_bit;
282    }
283
284    let half_exp = exponent - 127 + 15;
285    if half_exp >= 0x1f {
286        return sign | 0x7c00;
287    }
288    if half_exp <= 0 {
289        if half_exp < -10 {
290            return sign;
291        }
292        let mant = mantissa | 0x0080_0000;
293        let shift = (14 - half_exp) as u32;
294        let mut half_mant = (mant >> shift) as u16;
295        if ((mant >> (shift - 1)) & 1) != 0 {
296            half_mant = half_mant.saturating_add(1);
297        }
298        return sign | half_mant;
299    }
300
301    let mut half = sign | ((half_exp as u16) << 10) | ((mantissa >> 13) as u16);
302    if (mantissa & 0x0000_1000) != 0 {
303        half = half.saturating_add(1);
304    }
305    half
306}
skippy_protocol/binary/activation.rs

skippy_protocol/binary/
activation.rs
