File size: 6,640 Bytes
9425aed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
-- =====================================================================
-- SOV SCHEDULER -- C-- Continuous Batching State Machine
-- States: IDLE(0) PREFILL(1) GENERATE(2) SWAP(3) CHECKPOINT(4) RESUME(5)
-- Janet array: 32 Word32 slots at known address for runtime config
-- WORM: every 64 generated tokens → worm_checkpoint() → Blake3+Ed25519 receipt
-- =====================================================================

section "data" {
    -- Janet array: 32 Word32 slots
    -- [0]  = pending_request_count
    -- [1]  = batch_size
    -- [2]  = tokens_generated (mod 64 counter)
    -- [3]  = current_seq_id
    -- [4]  = kv_blocks_used
    -- [5]  = power_state (0=active,1=suspend,2=resume,3=low_battery)
    -- [6]  = speculative_draft_len
    -- [7]  = bft_quorum_height
    -- [8..15] = worm_receipt_blake3 (8 x Word32 = 32 bytes)
    -- [16..31] = worm_receipt_ed25519 first 64 bytes
    scheduler_janet_array:
        bits32[32] {0,0,0,0, 0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0};
}

-- External C/Fortran functions
foreign import ccall "sov_kv_allocate_blocks"
    sov_kv_allocate_blocks :: Word64 -> Word32 -> IO ();

foreign import ccall "sov_cuda_flash_attention"
    sov_cuda_flash_attention :: Word32 -> Word32 -> Word64 -> Word64 -> Word64 -> Word64 -> Word64 -> Word64 -> Word32 -> Word32 -> IO ();

foreign import ccall "sov_worm_checkpoint"
    sov_worm_checkpoint :: Word64 -> IO ();

foreign import ccall "sov_worm_restore"
    sov_worm_restore :: Word64 -> IO ();

foreign import ccall "sov_kv_append_tokens"
    sov_kv_append_tokens :: Word64 -> Word64 -> Word64 -> Word64 -> Word32 -> Word32 -> IO ();

foreign import ccall "sov_speculative_draft"
    sov_speculative_draft :: Word64 -> Word32 -> Word64 -> IO ();

foreign import ccall "sov_speculative_verify"
    sov_speculative_verify :: Word64 -> Word64 -> Word32 -> IO Word32;

foreign import ccall "sov_power_event"
    sov_power_event :: Word32 -> IO ();

-- =====================================================================
-- scheduler_step(state, batch_ptr, kv_ptr) -> new_state
-- =====================================================================
scheduler_step(W_ state, W_ batch_ptr, W_ kv_ptr)
{
    W_ pending, new_state, tokens_gen, power;

    power = W_[scheduler_janet_array + (5 * SIZEOF_W)];
    if (power == 1) {   -- SUSPEND
        jump scheduler_do_checkpoint(state, batch_ptr, kv_ptr);
    }
    if (power == 2) {   -- RESUME
        jump scheduler_do_resume(state, batch_ptr, kv_ptr);
    }

    switch [0..5] state {
        case 0: {   -- IDLE
            pending = W_[scheduler_janet_array + (0 * SIZEOF_W)];
            if (pending > 0) {
                new_state = 1;  -- -> PREFILL
            } else {
                new_state = 0;  -- stay IDLE
            }
            return (new_state);
        }
        case 1: {   -- PREFILL
            -- Allocate KV blocks for this batch
            foreign "C" sov_kv_allocate_blocks(kv_ptr, W_[scheduler_janet_array + (1 * SIZEOF_W)]);
            -- Reset token counter
            W_[scheduler_janet_array + (2 * SIZEOF_W)] = 0;
            new_state = 2;  -- -> GENERATE
            return (new_state);
        }
        case 2: {   -- GENERATE
            W_ heads, seqs, head_dim, block_size;
            seqs      = W_[scheduler_janet_array + (1 * SIZEOF_W)];
            heads     = 32;    -- llama-3 8B default
            head_dim  = 128;
            block_size = 16;

            -- Run flash attention forward pass
            foreign "C" sov_cuda_flash_attention(
                seqs, heads,
                batch_ptr,          -- q
                kv_ptr,             -- k
                kv_ptr + 4096,      -- v
                batch_ptr + 8192,   -- out
                kv_ptr + 16384,     -- block_table
                kv_ptr + 32768,     -- seq_lens
                head_dim, block_size);

            -- Run speculative draft
            foreign "C" sov_speculative_draft(batch_ptr, W_[scheduler_janet_array + (6 * SIZEOF_W)], kv_ptr);

            -- Increment token counter
            tokens_gen = W_[scheduler_janet_array + (2 * SIZEOF_W)] + 1;
            W_[scheduler_janet_array + (2 * SIZEOF_W)] = tokens_gen;

            -- Every 64 tokens: WORM checkpoint
            if ((tokens_gen & 63) == 0) {
                foreign "C" sov_worm_checkpoint(kv_ptr);
                -- BFT vote: increment quorum height
                W_[scheduler_janet_array + (7 * SIZEOF_W)] = W_[scheduler_janet_array + (7 * SIZEOF_W)] + 1;
            }

            -- Check if sequence complete (token_count reached max or EOS)
            if (tokens_gen >= 2048) {
                -- Decrement pending count
                W_[scheduler_janet_array + (0 * SIZEOF_W)] = W_[scheduler_janet_array + (0 * SIZEOF_W)] - 1;
                new_state = 0;  -- -> IDLE
            } else {
                new_state = 2;  -- stay GENERATE
            }
            return (new_state);
        }
        case 3: {   -- SWAP: copy KV blocks GPU->CPU
            -- Swap oldest sequence KV to CPU (simplified: just checkpoint)
            foreign "C" sov_worm_checkpoint(kv_ptr);
            new_state = 1;  -- -> PREFILL
            return (new_state);
        }
        case 4: {   -- CHECKPOINT
            jump scheduler_do_checkpoint(state, batch_ptr, kv_ptr);
        }
        case 5: {   -- RESUME
            jump scheduler_do_resume(state, batch_ptr, kv_ptr);
        }
        default: {
            new_state = 0;
            return (new_state);
        }
    }
}

scheduler_do_checkpoint(W_ state, W_ batch_ptr, W_ kv_ptr)
{
    foreign "C" sov_worm_checkpoint(kv_ptr);
    -- Receipt written to janet_array[8..31] by worm_checkpoint C impl
    W_[scheduler_janet_array + (5 * SIZEOF_W)] = 0;  -- clear power flag
    return (5);  -- -> RESUME state
}

scheduler_do_resume(W_ state, W_ batch_ptr, W_ kv_ptr)
{
    foreign "C" sov_worm_restore(kv_ptr);
    W_[scheduler_janet_array + (5 * SIZEOF_W)] = 0;  -- clear power flag
    return (1);  -- -> PREFILL
}

-- =====================================================================
-- janet_get(slot) / janet_set(slot, val)
-- =====================================================================
janet_get(W_ slot)
{
    W_ val;
    val = W_[scheduler_janet_array + (slot * SIZEOF_W)];
    return (val);
}

janet_set(W_ slot, W_ val)
{
    W_[scheduler_janet_array + (slot * SIZEOF_W)] = val;
    return ();
}