superfluid_worker/
materialize.rs

1//! Wire ⇄ ABI materialization.
2
3use superfluid_abi::{
4    array::read_array, ArenaBounds, AdmitResult, LaneAdmit, LaneCommit,
5    LaneDecode, LaneEmit, LaneFault, LanePrefill, LaneRetire,
6    MatchCandidate, MatchResult, OpComplete, RingRef, SamplingParams,
7    ShedEntry, ShedPolicy, SpaceMatch, StateSpaceDesc, TickEvents,
8    TickPlan, TokenRef, RecordArena,
9};
10use superfluid_proto::handshake::StateSpaceDescMsg;
11use superfluid_proto::linkw::*;
12
13pub fn token_ref_from_msg(m: &TokenRefMsg) -> TokenRef {
14    TokenRef {
15        ring_id: m.ring_id,
16        index: m.index,
17        count: m.count,
18        _pad0: 0,
19        generation: m.generation,
20    }
21}
22
23pub fn token_ref_to_msg(r: &TokenRef) -> TokenRefMsg {
24    TokenRefMsg {
25        ring_id: r.ring_id,
26        index: r.index,
27        count: r.count,
28        generation: r.generation,
29    }
30}
31
32pub fn ring_ref_to_msg(r: &RingRef) -> RingRefMsg {
33    RingRefMsg {
34        ring_id: r.ring_id,
35        index: r.index,
36        generation: r.generation,
37    }
38}
39
40fn sampling_params(m: &SamplingParamsMsg) -> SamplingParams {
41    SamplingParams {
42        temperature: m.temperature,
43        top_p: m.top_p,
44        min_p: m.min_p,
45        top_k: m.top_k,
46        freq_penalty: m.freq_penalty,
47        presence_penalty: m.presence_penalty,
48        repeat_penalty: m.repeat_penalty,
49        flags: m.flags,
50    }
51}
52
53pub fn materialize_plan(msg: &TickPlanMsg) -> (TickPlan, RecordArena) {
54    let mut arena = RecordArena::new();
55
56    let admits: Vec<LaneAdmit> = msg
57        .admits
58        .iter()
59        .map(|a| LaneAdmit {
60            lane_tag: a.lane_tag,
61            prompt: token_ref_from_msg(&a.prompt),
62            seed_handle: a.seed_handle,
63            sampling: a.sampling,
64            logit_bias_handle: a.logit_bias_handle,
65            params: sampling_params(&a.params),
66            rng_counter_base: a.rng_counter_base,
67            grammar_handle: a.grammar_handle,
68            strategy_slot: a.strategy_slot,
69            determinism_class: a.determinism_class,
70            minimum_exactness: a.minimum_exactness,
71            allow_approximate: a.allow_approximate as u8,
72            want_logprobs: if a.want_logprobs { 1 + a.top_logprobs.min(20) } else { 0 },
73            required_cert_id: a.required_cert_id,
74            host_sampler_identity: a.host_sampler_identity,
75            grammar_replay: a.grammar_replay,
76            decode_replay: a.decode_replay,
77        })
78        .collect();
79    let commits: Vec<LaneCommit> = msg
80        .commits
81        .iter()
82        .map(|c| LaneCommit {
83            lane_tag: c.lane_tag,
84            token_id: c.token_id,
85            _pad0: 0,
86            logits_nonce: c.logits_nonce,
87        })
88        .collect();
89    let prefills: Vec<LanePrefill> = msg
90        .prefills
91        .iter()
92        .map(|p| LanePrefill {
93            lane_tag: p.lane_tag,
94            token_offset: p.token_offset,
95            token_count: p.token_count,
96        })
97        .collect();
98    let decodes: Vec<LaneDecode> = msg
99        .decodes
100        .iter()
101        .map(|d| LaneDecode {
102            lane_tag: d.lane_tag,
103            max_new_tokens: d.max_new_tokens,
104            overshoot: d.overshoot,
105            _pad1: 0,
106        })
107        .collect();
108    let retires: Vec<LaneRetire> = msg
109        .retires
110        .iter()
111        .map(|r| LaneRetire {
112            lane_tag: r.lane_tag,
113            publish_to_cache: r.publish_to_cache as u8,
114            _pad0: [0; 7],
115        })
116        .collect();
117
118    let plan = TickPlan {
119        struct_size: std::mem::size_of::<TickPlan>() as u64,
120        plan_seq: msg.plan_seq,
121        flags: msg.flags,
122        _pad0: 0,
123        prefill_token_budget: msg.prefill_token_budget,
124        max_decode_lanes: msg.max_decode_lanes,
125        admits: arena.push_records(&admits),
126        commits: arena.push_records(&commits),
127        prefills: arena.push_records(&prefills),
128        decodes: arena.push_records(&decodes),
129        retires: arena.push_records(&retires),
130        shed_policy: ShedPolicy {
131            victim_lanes: arena.push_records(&msg.shed_policy.victim_lanes),
132            evictable_cache_classes: msg.shed_policy.evictable_cache_classes,
133            protected_quota_bytes: msg.shed_policy.protected_quota_bytes,
134            max_evict_bytes: msg.shed_policy.max_evict_bytes,
135        },
136        on_partial_emit: None,
137        partial_emit_user: std::ptr::null_mut(),
138    };
139    (plan, arena)
140}
141
142pub fn events_to_msg(ev: &TickEvents) -> TickEventsMsg {
143    let bounds = ArenaBounds {
144        base: 0,
145        len: usize::MAX,
146    };
147    // SAFETY: engine-owned arrays valid until the next tick on this
148    // bundle; serialization happens before any next tick by construction.
149    unsafe {
150        TickEventsMsg {
151            plan_seq: ev.plan_seq,
152            tick_status: ev.tick_status,
153            admit_results: read_array::<AdmitResult>(&ev.admit_results, &bounds)
154                .map(|it| {
155                    it.map(|r| AdmitResultMsg {
156                        lane_tag: r.lane_tag,
157                        status: r.status,
158                        reject_code: r.reject_code,
159                        cert_id: r.cert_id,
160                        granted_class: r.granted_class,
161                    })
162                    .collect()
163                })
164                .unwrap_or_default(),
165            emits: read_array::<LaneEmit>(&ev.emits, &bounds)
166                .map(|it| {
167                    it.map(|e| LaneEmitMsg {
168                        lane_tag: e.lane_tag,
169                        token_ref: token_ref_to_msg(&e.token_ref),
170                        n_tokens: e.n_tokens,
171                        finish: e.finish,
172                        logits_row: ring_ref_to_msg(&e.logits_row),
173                        spec: SpecStatsMsg {
174                            proposed: e.spec.proposed,
175                            accepted: e.spec.accepted,
176                            cert_id_in_effect: e.spec.cert_id_in_effect,
177                        },
178                    })
179                    .collect()
180                })
181                .unwrap_or_default(),
182            faults: read_array::<LaneFault>(&ev.faults, &bounds)
183                .map(|it| {
184                    it.map(|f| LaneFaultMsg {
185                        lane_tag: f.lane_tag,
186                        code: f.code,
187                        detail: f.detail,
188                    })
189                    .collect()
190                })
191                .unwrap_or_default(),
192            shed: ShedReportMsg {
193                prefill_chunks_dropped: ev.shed.prefill_chunks_dropped,
194                evictions_performed: ev.shed.evictions_performed,
195                bytes_evicted: ev.shed.bytes_evicted,
196                entries: read_array::<ShedEntry>(&ev.shed.entries, &bounds)
197                    .map(|it| {
198                        it.map(|s| ShedEntryMsg {
199                            lane_tag: s.lane_tag,
200                            kind: s.kind,
201                            reason: s.reason,
202                            bytes: s.bytes,
203                            tokens: s.tokens,
204                        })
205                        .collect()
206                    })
207                    .unwrap_or_default(),
208            },
209            op_completions: read_array::<OpComplete>(&ev.op_completions, &bounds)
210                .map(|it| {
211                    it.map(|o| OpCompleteMsg {
212                        op: o.op,
213                        state: o.state,
214                        error: o.error,
215                        bytes_moved: o.bytes_moved,
216                    })
217                    .collect()
218                })
219                .unwrap_or_default(),
220            timings: TickTimingsMsg {
221                wall_ns: ev.timings.wall_ns,
222                prefill_ns: ev.timings.prefill_ns,
223                decode_ns: ev.timings.decode_ns,
224                graph_hits: ev.timings.graph_hits,
225                graph_misses: ev.timings.graph_misses,
226            },
227            mem: MemCountersMsg {
228                allocated_bytes: ev.mem.allocated_bytes,
229                host_retained_bytes: ev.mem.host_retained_bytes,
230                pool_blocks_total: ev.mem.pool_blocks_total,
231                pool_blocks_used: ev.mem.pool_blocks_used,
232                pool_bytes_evictable: ev.mem.pool_bytes_evictable,
233                tentative_bytes: ev.mem.tentative_bytes,
234            },
235            logprobs: read_array::<superfluid_abi::LaneLogprob>(&ev.logprobs, &bounds)
236                .map(|it| {
237                    it.map(|l| {
238                        let n = (l.n_top as usize).min(superfluid_abi::MAX_TOP_LOGPROBS);
239                        superfluid_proto::linkw::LaneLogprobMsg {
240                            lane_tag: l.lane_tag,
241                            logprob_bits: l.logprob.to_bits(),
242                            top_ids: l.top_ids[..n].to_vec(),
243                            top_logprob_bits: l.top_logprobs[..n].iter().map(|f| f.to_bits()).collect(),
244                        }
245                    })
246                    .collect()
247                })
248                .unwrap_or_default(),
249        }
250    }
251}
252
253pub fn desc_to_msg(d: &StateSpaceDesc) -> StateSpaceDescMsg {
254    let name_end = d.name.iter().position(|b| *b == 0).unwrap_or(d.name.len());
255    StateSpaceDescMsg {
256        space_id: d.space_id,
257        kind: d.kind,
258        version_tag: d.version_tag,
259        bytes_per_token: d.bytes_per_token,
260        blob_bytes: d.blob_bytes,
261        page_size_tokens: d.page_size_tokens,
262        fork_cost_class: d.fork_cost_class,
263        fork_cost_bytes: d.fork_cost_bytes,
264        snapshot_cadence: d.snapshot_cadence,
265        snapshot_interval_tokens: d.snapshot_interval_tokens,
266        placement: d.placement,
267        flags: d.flags,
268        name: String::from_utf8_lossy(&d.name[..name_end]).into_owned(),
269    }
270}
271
272pub fn match_result_to_msg(m: &MatchResult) -> MatchResMsg {
273    let bounds = ArenaBounds {
274        base: 0,
275        len: usize::MAX,
276    };
277    // SAFETY: engine-owned match result, next-call lifetime.
278    let spaces = unsafe {
279        read_array::<SpaceMatch>(&m.spaces, &bounds)
280            .map(|it| {
281                it.map(|s| SpaceMatchMsg {
282                    space_id: s.space_id,
283                    candidates: read_array::<MatchCandidate>(&s.candidates, &bounds)
284                        .map(|cands| {
285                            cands
286                                .map(|c| MatchCandidateMsg {
287                                    prefix_len: c.prefix_len,
288                                    provenance_digest: c.provenance_digest,
289                                    taint_bits: c.taint_bits,
290                                    resident_tier: c.resident_tier,
291                                })
292                                .collect()
293                        })
294                        .unwrap_or_default(),
295                })
296                .collect()
297            })
298            .unwrap_or_default()
299    };
300    MatchResMsg { spaces }
301}
superfluid_worker/materialize.rs

superfluid_worker/
materialize.rs
