skippy-server 0.77.0

Embedded Skippy staged runtime server
use skippy_runtime::{
    ACTIVATION_BOUNDARY_DESC_VERSION, ACTIVATION_FRAME_VERSION, ACTIVATION_IDENTITY_BYTES,
    ACTIVATION_MAX_PARTS, ActivationBoundaryDesc, ActivationDesc, ActivationFrame,
    ActivationPartDesc, GGML_TYPE_F32, GGML_TYPE_I32,
};

pub(crate) fn boundary_f32(elements_per_token: u64) -> ActivationBoundaryDesc {
    let mut parts = [ActivationPartDesc::default(); ACTIVATION_MAX_PARTS];
    parts[0] = ActivationPartDesc {
        identity: [1; ACTIVATION_IDENTITY_BYTES],
        ggml_type: GGML_TYPE_F32,
        rank: 2,
        token_axis: 1,
        dimensions: [elements_per_token as i64, -1, 0, 0],
        byte_strides: [4, elements_per_token.saturating_mul(4), 0, 0],
        ..ActivationPartDesc::default()
    };
    ActivationBoundaryDesc {
        version: ACTIVATION_BOUNDARY_DESC_VERSION,
        part_count: 1,
        frontier_identity: [9; ACTIVATION_IDENTITY_BYTES],
        parts,
    }
}

pub(crate) struct PartBytes {
    pub identity: u8,
    pub ggml_type: u32,
    pub flags: u32,
    pub bytes: Vec<u8>,
}

pub(crate) fn frame(token_count: u32, part_specs: Vec<PartBytes>) -> ActivationFrame {
    let mut payload = Vec::new();
    let mut parts = [ActivationPartDesc::default(); ACTIVATION_MAX_PARTS];
    for (index, spec) in part_specs.iter().enumerate() {
        let element_bytes = match spec.ggml_type {
            GGML_TYPE_F32 | GGML_TYPE_I32 => 4,
            other => panic!("unsupported test activation type {other}"),
        };
        let token_count_usize = token_count as usize;
        assert!(token_count_usize > 0);
        assert_eq!(spec.bytes.len() % token_count_usize, 0);
        let row_bytes = spec.bytes.len() / token_count_usize;
        assert_eq!(row_bytes % element_bytes, 0);
        parts[index] = ActivationPartDesc {
            identity: [spec.identity; ACTIVATION_IDENTITY_BYTES],
            ggml_type: spec.ggml_type,
            rank: 2,
            token_axis: 1,
            flags: spec.flags,
            dimensions: [
                (row_bytes / element_bytes) as i64,
                i64::from(token_count),
                0,
                0,
            ],
            byte_strides: [element_bytes as u64, row_bytes as u64, 0, 0],
            payload_offset: payload.len() as u64,
            payload_bytes: spec.bytes.len() as u64,
        };
        payload.extend_from_slice(&spec.bytes);
    }
    ActivationFrame {
        desc: ActivationDesc {
            version: ACTIVATION_FRAME_VERSION,
            producer_stage_index: 1,
            layer_start: 4,
            layer_end: 8,
            token_count,
            sequence_count: 1,
            part_count: part_specs.len() as u32,
            payload_bytes: payload.len() as u64,
            frontier_identity: [9; ACTIVATION_IDENTITY_BYTES],
            parts,
        },
        payload,
    }
}

pub(crate) fn f32_frame(token_count: u32, values: &[f32]) -> ActivationFrame {
    frame(
        token_count,
        vec![PartBytes {
            identity: 1,
            ggml_type: GGML_TYPE_F32,
            flags: 0,
            bytes: values
                .iter()
                .flat_map(|value| value.to_le_bytes())
                .collect(),
        }],
    )
}