pie 0.1.1

Programmable Inference Engine (PIE)
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

interface forward-text {

    use wasi:io/poll@0.2.4.{pollable};
    use core.{queue};
    use allocate.{object-id};

    // Executes a forward pass using context and input embeddings
    forward-text: func(
        queue: borrow<queue>,                   // Queue to schedule the operation
        last-kv-page-len: u32,                  // Number of valid entries in the last KV page
        kv-page-ids: list<object-id>,           // List of context key-value memory pages
        tokens: list<u32>,                      // List of input token IDs
        positions: list<u32>,                   // List of input token positions
        mask: list<list<u32>>,                  // 2d mask in BRLE format
        output-indices: list<u32>,              // Output indices
    ) -> distribution-result;

    forward-text-no-output: func(
        queue: borrow<queue>,                   // Queue to schedule the operation
        last-kv-page-len: u32,                  // Number of valid entries in the last KV page
        kv-page-ids: list<object-id>,           // List of context key-value memory pages
        tokens: list<u32>,                      // List of input token IDs
        positions: list<u32>,                   // List of input token positions
        mask: list<list<u32>>                   // 2d mask in BRLE format
    );

    resource distribution-result {

        // Returns a pollable object to check when the result is ready
        pollable: func() -> pollable;

        // Retrieves the result if ready; None if still pending
        // Each tuple: (token IDs, associated probabilities)
        get: func() -> option<list<tuple<list<u32>, list<f32>>>>;
    }
}