skippy_protocol/binary/
types.rs

1use std::io;
2
3use super::{
4    activation::{decode_f16_to_f32_bytes, decode_q8_to_f32_bytes_with_state_flags},
5    invalid_data,
6};
7
8// v11 adds the Inkling MTP embedding sideband and makes the coordinator the sole owner of
9// verify-window acceptance, removing redundant tail-stage acceptance/correction fields. Stage
10// peers must be upgraded together so older readers reject the changed payload contract.
11pub const STAGE_STATE_VERSION: i32 = 11;
12pub const MAX_STAGE_LOGIT_BIAS: usize = 256;
13pub const MAX_STAGE_PREDICTED_TOKENS: usize = 262_144;
14pub const MAX_STAGE_SIDEBAND_VALUES: usize = 1_048_576;
15pub const MAX_STAGE_CHAT_SAMPLING_METADATA_BYTES: usize = 8 * 1024 * 1024;
16pub const MAX_STAGE_STATE_IMPORT_BYTES: usize = 512 * 1024 * 1024;
17pub const MAX_STAGE_ACTIVATION_BYTES: usize = 512 * 1024 * 1024;
18pub const MAX_STAGE_DECODED_ACTIVATION_BYTES: usize = 512 * 1024 * 1024;
19pub const READY_MAGIC: i32 = 0x5352_4459; // "SRDY"
20pub const LLAMA_TOKEN_NULL: i32 = -1;
21pub const STAGE_STATE_HEADER_BYTES: usize = 10 * 4;
22pub const STAGE_SAMPLING_CONFIG_BASE_BYTES: usize = 10 * 4;
23pub const STAGE_LOGIT_BIAS_WIRE_BYTES: usize = 4 + 4;
24pub const STAGE_WIRE_FIXED_HEADER_BYTES: usize = 5 * 4 + STAGE_STATE_HEADER_BYTES + 2 * 8;
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27#[repr(i32)]
28pub enum WireActivationDType {
29    F32 = 0,
30    F16 = 1,
31    Q8 = 2,
32}
33
34impl TryFrom<i32> for WireActivationDType {
35    type Error = io::Error;
36
37    fn try_from(value: i32) -> Result<Self, Self::Error> {
38        match value {
39            0 => Ok(Self::F32),
40            1 => Ok(Self::F16),
41            2 => Ok(Self::Q8),
42            _ => Err(invalid_data("unknown activation wire dtype")),
43        }
44    }
45}
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48#[repr(i32)]
49pub enum WireMessageKind {
50    PrefillEmbd = 1,
51    DecodeEmbd = 2,
52    Stop = 3,
53    PrefillFinalEmbd = 4,
54    DecodeReplayEmbd = 5,
55    DecodeReplayFinalEmbd = 6,
56    StateImport = 7,
57    DecodeReadout = 8,
58    DecodeLightCtx = 9,
59    VerifyWindow = 21,
60    RetireVerifyWindow = 22,
61    StateExport = 13,
62    ConfigureGeneration = 14,
63    ProbePrefill = 15,
64    RestorePrefill = 16,
65    TryRestorePrefill = 17,
66    TryRestorePrefillDecode = 18,
67    TrimSession = 19,
68    PredictionReturnOpen = 20,
69}
70
71impl WireMessageKind {
72    pub fn is_prefill(self) -> bool {
73        matches!(self, Self::PrefillEmbd | Self::PrefillFinalEmbd)
74    }
75
76    pub fn is_decode_replay(self) -> bool {
77        matches!(self, Self::DecodeReplayEmbd | Self::DecodeReplayFinalEmbd)
78    }
79
80    pub fn is_decode_light_context(self) -> bool {
81        matches!(self, Self::DecodeLightCtx)
82    }
83
84    pub fn requires_predicted_reply(self) -> bool {
85        matches!(
86            self,
87            Self::DecodeEmbd
88                | Self::DecodeReadout
89                | Self::DecodeLightCtx
90                | Self::VerifyWindow
91                | Self::PrefillFinalEmbd
92                | Self::DecodeReplayFinalEmbd
93        )
94    }
95
96    pub fn is_session_control(self) -> bool {
97        matches!(self, Self::TrimSession)
98    }
99
100    pub fn is_verify_retirement(self) -> bool {
101        matches!(self, Self::RetireVerifyWindow)
102    }
103
104    pub fn is_generation_control(self) -> bool {
105        matches!(self, Self::ConfigureGeneration)
106    }
107
108    pub fn is_prefix_cache_control(self) -> bool {
109        matches!(
110            self,
111            Self::ProbePrefill
112                | Self::RestorePrefill
113                | Self::TryRestorePrefill
114                | Self::TryRestorePrefillDecode
115        )
116    }
117
118    pub fn is_activationless_prefix_cache_control(self) -> bool {
119        matches!(
120            self,
121            Self::ProbePrefill | Self::RestorePrefill | Self::TryRestorePrefill
122        )
123    }
124}
125
126impl TryFrom<i32> for WireMessageKind {
127    type Error = io::Error;
128
129    fn try_from(value: i32) -> Result<Self, Self::Error> {
130        match value {
131            1 => Ok(Self::PrefillEmbd),
132            2 => Ok(Self::DecodeEmbd),
133            3 => Ok(Self::Stop),
134            4 => Ok(Self::PrefillFinalEmbd),
135            5 => Ok(Self::DecodeReplayEmbd),
136            6 => Ok(Self::DecodeReplayFinalEmbd),
137            7 => Ok(Self::StateImport),
138            8 => Ok(Self::DecodeReadout),
139            9 => Ok(Self::DecodeLightCtx),
140            13 => Ok(Self::StateExport),
141            14 => Ok(Self::ConfigureGeneration),
142            15 => Ok(Self::ProbePrefill),
143            16 => Ok(Self::RestorePrefill),
144            17 => Ok(Self::TryRestorePrefill),
145            18 => Ok(Self::TryRestorePrefillDecode),
146            19 => Ok(Self::TrimSession),
147            20 => Ok(Self::PredictionReturnOpen),
148            21 => Ok(Self::VerifyWindow),
149            22 => Ok(Self::RetireVerifyWindow),
150            _ => Err(invalid_data("unknown stage message kind")),
151        }
152    }
153}
154
155#[derive(Debug, Clone, Copy, PartialEq, Eq)]
156#[repr(i32)]
157pub enum WireReplyKind {
158    Ack = 1,
159    PredictedToken = 2,
160    PredictedTokens = 3,
161}
162
163impl TryFrom<i32> for WireReplyKind {
164    type Error = io::Error;
165
166    fn try_from(value: i32) -> Result<Self, Self::Error> {
167        match value {
168            1 => Ok(Self::Ack),
169            2 => Ok(Self::PredictedToken),
170            3 => Ok(Self::PredictedTokens),
171            _ => Err(invalid_data("unknown stage reply kind")),
172        }
173    }
174}
175
176#[derive(Debug, Clone, Copy, PartialEq, Eq)]
177#[repr(i32)]
178pub enum WireStagePhase {
179    Prefill = 1,
180    Decode = 2,
181    DecodeReplay = 3,
182    DecodeLight = 4,
183}
184
185pub mod state_flags {
186    pub const FINAL_CHUNK: i32 = 1 << 0;
187    pub const LIGHT_CONTEXT: i32 = 1 << 1;
188    pub const SAMPLING: i32 = 1 << 3;
189    pub const FULL_STATE: i32 = 1 << 4;
190    pub const CHAT_SAMPLING_METADATA: i32 = 1 << 5;
191    pub const RWKV7_V_FIRST_SIDEBAND: i32 = 1 << 6;
192    pub const GEMMA3N_ALTUP_SIDEBAND: i32 = 1 << 7;
193    pub const INKLING_MTP_EMBD_SIDEBAND: i32 = 1 << 8;
194}
195
196pub const ACTIVATION_FLAG_RWKV7_V_FIRST: u64 = 1 << 0;
197pub const ACTIVATION_FLAG_GEMMA3N_ALTUP: u64 = 1 << 1;
198pub const ACTIVATION_FLAG_INKLING_MTP_EMBD: u64 = 1 << 2;
199
200pub fn activation_frame_flags_from_state_flags(flags: i32) -> u64 {
201    let mut frame_flags = 0;
202    if (flags & state_flags::RWKV7_V_FIRST_SIDEBAND) != 0 {
203        frame_flags |= ACTIVATION_FLAG_RWKV7_V_FIRST;
204    }
205    if (flags & state_flags::GEMMA3N_ALTUP_SIDEBAND) != 0 {
206        frame_flags |= ACTIVATION_FLAG_GEMMA3N_ALTUP;
207    }
208    if (flags & state_flags::INKLING_MTP_EMBD_SIDEBAND) != 0 {
209        frame_flags |= ACTIVATION_FLAG_INKLING_MTP_EMBD;
210    }
211    frame_flags
212}
213
214pub fn activation_state_flags_from_frame_flags(flags: u64) -> i32 {
215    let mut state = 0;
216    if (flags & ACTIVATION_FLAG_RWKV7_V_FIRST) != 0 {
217        state |= state_flags::RWKV7_V_FIRST_SIDEBAND;
218    }
219    if (flags & ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0 {
220        state |= state_flags::GEMMA3N_ALTUP_SIDEBAND;
221    }
222    if (flags & ACTIVATION_FLAG_INKLING_MTP_EMBD) != 0 {
223        state |= state_flags::INKLING_MTP_EMBD_SIDEBAND;
224    }
225    state
226}
227
228#[derive(Debug, Clone, Copy, PartialEq, Eq)]
229pub struct StageStateHeader {
230    pub version: i32,
231    pub seq_id: i32,
232    pub phase: i32,
233    pub flags: i32,
234    pub checkpoint_generation: i32,
235    pub prompt_token_count: i32,
236    pub decode_step: i32,
237    pub current_token: i32,
238    pub source_stage_index: i32,
239    pub reserved: i32,
240}
241
242#[derive(Debug, Clone, Copy, PartialEq)]
243pub struct StageLogitBias {
244    pub token_id: i32,
245    pub bias: f32,
246}
247
248impl Default for StageLogitBias {
249    fn default() -> Self {
250        Self {
251            token_id: 0,
252            bias: 0.0,
253        }
254    }
255}
256
257#[derive(Debug, Clone, PartialEq)]
258pub struct StageSamplingConfig {
259    pub flags: u32,
260    pub seed: u32,
261    pub temperature: f32,
262    pub top_p: f32,
263    pub top_k: i32,
264    pub min_p: f32,
265    pub presence_penalty: f32,
266    pub frequency_penalty: f32,
267    pub repeat_penalty: f32,
268    pub penalty_last_n: i32,
269    pub logit_bias: Vec<StageLogitBias>,
270}
271
272impl Default for StageSamplingConfig {
273    fn default() -> Self {
274        Self {
275            flags: 0,
276            seed: 0,
277            temperature: 1.0,
278            top_p: 1.0,
279            top_k: 0,
280            min_p: 0.0,
281            presence_penalty: 0.0,
282            frequency_penalty: 0.0,
283            repeat_penalty: 1.0,
284            penalty_last_n: -1,
285            logit_bias: Vec::new(),
286        }
287    }
288}
289
290impl StageSamplingConfig {
291    pub fn enabled(&self) -> bool {
292        self.flags != 0
293    }
294}
295
296impl StageStateHeader {
297    pub fn new(kind: WireMessageKind, dtype: WireActivationDType) -> Self {
298        let mut header = Self {
299            version: STAGE_STATE_VERSION,
300            seq_id: 0,
301            phase: expected_phase(kind) as i32,
302            flags: 0,
303            checkpoint_generation: 0,
304            prompt_token_count: 0,
305            decode_step: -1,
306            current_token: LLAMA_TOKEN_NULL,
307            source_stage_index: -1,
308            reserved: dtype as i32,
309        };
310        if matches!(
311            kind,
312            WireMessageKind::PrefillFinalEmbd | WireMessageKind::DecodeReplayFinalEmbd
313        ) {
314            header.flags |= state_flags::FINAL_CHUNK;
315        }
316        if kind.is_decode_light_context() {
317            header.flags |= state_flags::LIGHT_CONTEXT;
318        }
319        header
320    }
321
322    pub fn dtype(self) -> io::Result<WireActivationDType> {
323        WireActivationDType::try_from(self.reserved)
324    }
325
326    pub fn matches_kind(self, kind: WireMessageKind) -> bool {
327        if matches!(
328            kind,
329            WireMessageKind::StateImport | WireMessageKind::StateExport
330        ) || kind.is_session_control()
331            || kind.is_verify_retirement()
332            || kind.is_generation_control()
333        {
334            return true;
335        }
336        if self.phase != expected_phase(kind) as i32 {
337            return false;
338        }
339        let expected_final = matches!(
340            kind,
341            WireMessageKind::PrefillFinalEmbd | WireMessageKind::DecodeReplayFinalEmbd
342        );
343        let actual_final = (self.flags & state_flags::FINAL_CHUNK) != 0;
344        let expected_light = kind.is_decode_light_context();
345        let actual_light = (self.flags & state_flags::LIGHT_CONTEXT) != 0;
346        expected_final == actual_final && expected_light == actual_light
347    }
348}
349
350impl Default for StageStateHeader {
351    fn default() -> Self {
352        Self::new(WireMessageKind::PrefillEmbd, WireActivationDType::F32)
353    }
354}
355
356#[derive(Debug, Clone, Copy, PartialEq, Eq)]
357pub struct StageRequestEpoch {
358    pub request_id: u64,
359    pub session_id: u64,
360    pub checkpoint_generation: i32,
361    pub prompt_token_count: i32,
362    pub decode_step: i32,
363}
364
365impl StageRequestEpoch {
366    pub fn same_flow(self, other: Self) -> bool {
367        self.request_id == other.request_id && self.session_id == other.session_id
368    }
369
370    /// Returns true when this epoch is strictly older than `current` within the
371    /// same request/session flow.
372    ///
373    /// Epochs from different flows are never comparable. For matching flows,
374    /// staleness uses lexicographic ordering of checkpoint generation, prompt
375    /// token count, and decode step, so a newer checkpoint dominates prompt and
376    /// decode progress, and prompt progress dominates decode progress.
377    pub fn is_stale_for(self, current: Self) -> bool {
378        self.same_flow(current)
379            && (
380                self.checkpoint_generation,
381                self.prompt_token_count,
382                self.decode_step,
383            ) < (
384                current.checkpoint_generation,
385                current.prompt_token_count,
386                current.decode_step,
387            )
388    }
389}
390
391#[derive(Debug, Clone, PartialEq)]
392pub struct StageWireMessage {
393    pub kind: WireMessageKind,
394    pub pos_start: i32,
395    pub token_count: i32,
396    pub state: StageStateHeader,
397    pub request_id: u64,
398    pub session_id: u64,
399    pub sampling: Option<StageSamplingConfig>,
400    pub chat_sampling_metadata: Option<String>,
401    pub tokens: Vec<i32>,
402    pub positions: Vec<i32>,
403    pub activation: Vec<u8>,
404    pub raw_bytes: Vec<u8>,
405}
406
407impl StageWireMessage {
408    /// The committed session position that must exist before this message runs.
409    ///
410    /// Stage-state v11 makes this absolute position authoritative: a worker whose
411    /// speculative KV is ahead must rewind locally before executing the message.
412    pub fn authoritative_session_position(&self) -> Option<u64> {
413        if !matches!(
414            self.kind,
415            WireMessageKind::DecodeEmbd
416                | WireMessageKind::DecodeReadout
417                | WireMessageKind::DecodeLightCtx
418                | WireMessageKind::VerifyWindow
419        ) {
420            return None;
421        }
422        u64::try_from(self.pos_start).ok()
423    }
424
425    pub fn verify_window_id(&self) -> Option<i32> {
426        (self.kind == WireMessageKind::VerifyWindow).then_some(self.state.seq_id)
427    }
428
429    pub fn verify_window_base_position(&self) -> Option<i32> {
430        (self.kind == WireMessageKind::VerifyWindow).then_some(self.pos_start)
431    }
432
433    pub fn verify_window_token_count(&self) -> Option<i32> {
434        (self.kind == WireMessageKind::VerifyWindow).then_some(self.token_count)
435    }
436
437    pub fn estimated_wire_bytes(&self) -> usize {
438        let sampling_bytes = self.sampling.as_ref().map_or(0, |sampling| {
439            STAGE_SAMPLING_CONFIG_BASE_BYTES
440                + sampling.logit_bias.len().min(MAX_STAGE_LOGIT_BIAS) * STAGE_LOGIT_BIAS_WIRE_BYTES
441        });
442        let chat_metadata_bytes = self
443            .chat_sampling_metadata
444            .as_ref()
445            .map_or(0, |metadata| std::mem::size_of::<u32>() + metadata.len());
446        STAGE_WIRE_FIXED_HEADER_BYTES
447            .saturating_add(sampling_bytes)
448            .saturating_add(chat_metadata_bytes)
449            .saturating_add(self.payload_wire_bytes())
450    }
451
452    fn payload_wire_bytes(&self) -> usize {
453        if self.kind == WireMessageKind::StateImport {
454            return self.raw_bytes.len();
455        }
456        self.tokens
457            .len()
458            .saturating_mul(std::mem::size_of::<i32>())
459            .saturating_add(
460                self.positions
461                    .len()
462                    .saturating_mul(std::mem::size_of::<i32>()),
463            )
464            .saturating_add(self.activation.len())
465    }
466
467    pub fn request_epoch(&self) -> StageRequestEpoch {
468        StageRequestEpoch {
469            request_id: self.request_id,
470            session_id: self.session_id,
471            checkpoint_generation: self.state.checkpoint_generation,
472            prompt_token_count: self.state.prompt_token_count,
473            decode_step: self.state.decode_step,
474        }
475    }
476
477    pub fn stop(dtype: WireActivationDType) -> Self {
478        Self::stop_with_identity(dtype, 0, 0)
479    }
480
481    pub fn stop_with_identity(
482        dtype: WireActivationDType,
483        request_id: u64,
484        session_id: u64,
485    ) -> Self {
486        Self {
487            kind: WireMessageKind::Stop,
488            pos_start: 0,
489            token_count: 0,
490            state: StageStateHeader::new(WireMessageKind::Stop, dtype),
491            request_id,
492            session_id,
493            sampling: None,
494            chat_sampling_metadata: None,
495            tokens: Vec::new(),
496            positions: Vec::new(),
497            activation: Vec::new(),
498            raw_bytes: Vec::new(),
499        }
500    }
501
502    pub fn configure_generation(
503        dtype: WireActivationDType,
504        request_id: u64,
505        session_id: u64,
506        prompt_token_count: i32,
507        sampling: Option<StageSamplingConfig>,
508        chat_sampling_metadata: Option<String>,
509    ) -> Self {
510        let mut state = StageStateHeader::new(WireMessageKind::ConfigureGeneration, dtype);
511        state.prompt_token_count = prompt_token_count;
512        Self {
513            kind: WireMessageKind::ConfigureGeneration,
514            pos_start: 0,
515            token_count: 0,
516            state,
517            request_id,
518            session_id,
519            sampling,
520            chat_sampling_metadata,
521            tokens: Vec::new(),
522            positions: Vec::new(),
523            activation: Vec::new(),
524            raw_bytes: Vec::new(),
525        }
526    }
527
528    pub fn activation_f32_payload(&self, n_embd: i32) -> io::Result<Vec<u8>> {
529        if self.activation.is_empty() {
530            return Ok(Vec::new());
531        }
532        match self.state.dtype()? {
533            WireActivationDType::F32 => {
534                if self.activation.len() > MAX_STAGE_DECODED_ACTIVATION_BYTES {
535                    return Err(invalid_data(
536                        "decoded activation payload byte count exceeds maximum",
537                    ));
538                }
539                Ok(self.activation.clone())
540            }
541            WireActivationDType::F16 => decode_f16_to_f32_bytes(&self.activation),
542            WireActivationDType::Q8 => decode_q8_to_f32_bytes_with_state_flags(
543                &self.activation,
544                self.token_count,
545                n_embd,
546                self.state.flags,
547            ),
548        }
549    }
550
551    pub fn take_activation_f32_payload(&mut self, n_embd: i32) -> io::Result<Vec<u8>> {
552        if self.activation.is_empty() {
553            return Ok(Vec::new());
554        }
555        match self.state.dtype()? {
556            WireActivationDType::F32 => {
557                if self.activation.len() > MAX_STAGE_DECODED_ACTIVATION_BYTES {
558                    return Err(invalid_data(
559                        "decoded activation payload byte count exceeds maximum",
560                    ));
561                }
562                Ok(std::mem::take(&mut self.activation))
563            }
564            WireActivationDType::F16 => decode_f16_to_f32_bytes(&self.activation),
565            WireActivationDType::Q8 => decode_q8_to_f32_bytes_with_state_flags(
566                &self.activation,
567                self.token_count,
568                n_embd,
569                self.state.flags,
570            ),
571        }
572    }
573}
574
575#[derive(Debug, Clone, PartialEq, Eq)]
576pub struct StageReply {
577    pub kind: WireReplyKind,
578    pub predicted: i32,
579    /// Target-model predictions only. Native MTP proposals are carried separately.
580    pub predicted_tokens: Vec<i32>,
581    pub native_mtp_draft: Option<StageNativeMtpDraft>,
582    pub window: StageReplyWindow,
583    pub stats: StageReplyStats,
584}
585
586/// A native MTP proposal associated with a stage reply.
587///
588/// This deliberately has its own reply field rather than sharing the target
589/// prediction vector. Consumers must never mistake proposal metadata for a
590/// target prediction while verifying a composite speculative window.
591#[derive(Debug, Clone, PartialEq, Eq)]
592pub struct StageNativeMtpDraft {
593    pub token_ids: Vec<i32>,
594    pub proposal_compute_us: i64,
595}
596
597#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
598pub struct StageReplyWindow {
599    pub window_id: i32,
600}
601
602#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
603pub struct StageReplyStats {
604    pub kv_lookup_hits: i64,
605    pub kv_lookup_misses: i64,
606    pub kv_lookup_errors: i64,
607    pub kv_imported_pages: i64,
608    pub kv_imported_tokens: i64,
609    pub kv_recorded_pages: i64,
610    pub kv_recorded_bytes: i64,
611    pub kv_hit_stage_mask: i64,
612    pub kv_record_stage_mask: i64,
613    pub verify_window_compute_us: i64,
614    pub verify_window_forward_write_us: i64,
615    pub verify_window_downstream_wait_us: i64,
616    pub verify_window_total_us: i64,
617    pub verify_window_stage_count: i64,
618    pub verify_window_request_count: i64,
619    pub verify_window_token_count: i64,
620    pub verify_window_max_tokens: i64,
621    pub prefill_edge_write_us_max: i64,
622    pub prefill_edge_wait_us_max: i64,
623    pub prefill_edge_total_us_max: i64,
624    pub prefill_edge_stage_index: i64,
625    pub prefill_edge_activation_bytes_max: i64,
626    pub prefill_edge_observation_count: i64,
627}
628
629impl StageReplyStats {
630    pub fn merge(&mut self, other: Self) {
631        self.kv_lookup_hits += other.kv_lookup_hits;
632        self.kv_lookup_misses += other.kv_lookup_misses;
633        self.kv_lookup_errors += other.kv_lookup_errors;
634        self.kv_imported_pages += other.kv_imported_pages;
635        self.kv_imported_tokens += other.kv_imported_tokens;
636        self.kv_recorded_pages += other.kv_recorded_pages;
637        self.kv_recorded_bytes += other.kv_recorded_bytes;
638        self.kv_hit_stage_mask |= other.kv_hit_stage_mask;
639        self.kv_record_stage_mask |= other.kv_record_stage_mask;
640        self.verify_window_compute_us += other.verify_window_compute_us;
641        self.verify_window_forward_write_us += other.verify_window_forward_write_us;
642        self.verify_window_downstream_wait_us += other.verify_window_downstream_wait_us;
643        self.verify_window_total_us += other.verify_window_total_us;
644        self.verify_window_stage_count += other.verify_window_stage_count;
645        self.verify_window_request_count += other.verify_window_request_count;
646        self.verify_window_token_count += other.verify_window_token_count;
647        self.verify_window_max_tokens = self
648            .verify_window_max_tokens
649            .max(other.verify_window_max_tokens);
650        self.prefill_edge_write_us_max = self
651            .prefill_edge_write_us_max
652            .max(other.prefill_edge_write_us_max);
653        self.prefill_edge_wait_us_max = self
654            .prefill_edge_wait_us_max
655            .max(other.prefill_edge_wait_us_max);
656        if other.prefill_edge_total_us_max > self.prefill_edge_total_us_max {
657            self.prefill_edge_total_us_max = other.prefill_edge_total_us_max;
658            self.prefill_edge_stage_index = other.prefill_edge_stage_index;
659            self.prefill_edge_activation_bytes_max = other.prefill_edge_activation_bytes_max;
660        }
661        self.prefill_edge_observation_count += other.prefill_edge_observation_count;
662    }
663
664    pub fn observe_prefill_edge_transport(
665        &mut self,
666        stage_index: u32,
667        write_us: i64,
668        wait_us: i64,
669        activation_bytes: usize,
670    ) {
671        let write_us = write_us.max(0);
672        let wait_us = wait_us.max(0);
673        let total_us = write_us.saturating_add(wait_us);
674        self.prefill_edge_write_us_max = self.prefill_edge_write_us_max.max(write_us);
675        self.prefill_edge_wait_us_max = self.prefill_edge_wait_us_max.max(wait_us);
676        if total_us > self.prefill_edge_total_us_max {
677            self.prefill_edge_total_us_max = total_us;
678            self.prefill_edge_stage_index = i64::from(stage_index);
679            self.prefill_edge_activation_bytes_max =
680                i64::try_from(activation_bytes).unwrap_or(i64::MAX);
681        }
682        self.prefill_edge_observation_count = self.prefill_edge_observation_count.saturating_add(1);
683    }
684
685    pub fn is_empty(self) -> bool {
686        self.kv_lookup_hits == 0
687            && self.kv_lookup_misses == 0
688            && self.kv_lookup_errors == 0
689            && self.kv_imported_pages == 0
690            && self.kv_imported_tokens == 0
691            && self.kv_recorded_pages == 0
692            && self.kv_recorded_bytes == 0
693            && self.kv_hit_stage_mask == 0
694            && self.kv_record_stage_mask == 0
695            && self.verify_window_compute_us == 0
696            && self.verify_window_forward_write_us == 0
697            && self.verify_window_downstream_wait_us == 0
698            && self.verify_window_total_us == 0
699            && self.verify_window_stage_count == 0
700            && self.verify_window_request_count == 0
701            && self.verify_window_token_count == 0
702            && self.verify_window_max_tokens == 0
703            && self.prefill_edge_observation_count == 0
704    }
705}
706
707fn expected_phase(kind: WireMessageKind) -> WireStagePhase {
708    if kind.is_prefill()
709        || matches!(
710            kind,
711            WireMessageKind::StateImport | WireMessageKind::StateExport
712        )
713        || kind.is_session_control()
714        || kind.is_verify_retirement()
715        || kind.is_generation_control()
716        || kind.is_prefix_cache_control()
717    {
718        WireStagePhase::Prefill
719    } else if kind.is_decode_replay() {
720        WireStagePhase::DecodeReplay
721    } else if kind.is_decode_light_context() {
722        WireStagePhase::DecodeLight
723    } else {
724        WireStagePhase::Decode
725    }
726}
skippy_protocol/binary/types.rs

skippy_protocol/binary/
types.rs
