Skip to main content

ferrum_interfaces/
model_executor.rs

1//! Model execution interface with clear prefill/decode separation
2//!
3//! This module provides the ModelExecutor trait that replaces the "fat" Model
4//! interface, focusing purely on tensor operations without tokenization or sampling.
5
6use crate::{KvCacheHandle, RecurrentStateHandle, RecurrentStateSpec, TensorRef};
7use async_trait::async_trait;
8use ferrum_types::{ExecutorAdmissionLimits, FerrumError, ModelInfo, RequestId, Result, TokenId};
9use serde::{Deserialize, Serialize};
10use std::{
11    collections::{hash_map::DefaultHasher, HashMap, HashSet},
12    future::Future,
13    hash::{Hash, Hasher},
14    num::NonZeroU64,
15    ops::Range,
16    pin::Pin,
17    sync::Arc,
18};
19
20mod prefix_capture;
21mod prefix_restore;
22pub use prefix_capture::{
23    PrefixCaptureBoundary, PrefixCaptureLease, PrefixCapturePlan, PrefixCaptureRequest,
24    PrefixCaptureStatus,
25};
26pub use prefix_restore::{
27    PlanRuntimePrefixRestoreDeferral, PlanRuntimePrefixRestoreInput,
28    PlanRuntimePrefixRestoreOutcome, PlanRuntimePrefixRestoreOutput, PrefixRestoreDecision,
29    PrefixRestoreObservation, PrefixRestoreSource,
30};
31
32/// One model-owned KV slot reservation request.
33///
34/// `cache_id` is the executor/model cache key attached to a sequence. `target_len`
35/// is the sequence length that must be writable before the next forward runs.
36/// `admission_target_len`, when present, is a larger known-context bound used
37/// only for admission fit checks. Paged models must not allocate future blocks
38/// for it; it mirrors vLLM's chunked-prefill full-context fit gate.
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40pub struct KvSlotRequest {
41    pub cache_id: String,
42    pub target_len: usize,
43    pub admission_target_len: Option<usize>,
44}
45
46/// Per-cache outcome from a KV slot reservation attempt.
47#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
48pub struct KvSlotAllocation {
49    pub cache_id: String,
50    pub blocks_before: usize,
51    pub blocks_after: usize,
52    pub new_blocks: usize,
53}
54
55/// Model-owned paged-KV reservation evidence.
56///
57/// Executors that own a vLLM-style physical KV block pool return this after
58/// reserving all requested slots. Executors without model-owned paged KV return
59/// `None` from `reserve_kv_slots`.
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61pub struct KvSlotReservation {
62    pub block_size: usize,
63    pub total_blocks: usize,
64    pub free_blocks_before: usize,
65    pub free_blocks_after: usize,
66    pub allocations: Vec<KvSlotAllocation>,
67}
68
69/// Point-in-time model-owned paged-KV capacity snapshot.
70///
71/// This is intentionally smaller than [`KvSlotReservation`]: it lets the
72/// engine observe whether physical block capacity has actually changed after a
73/// release, without allocating speculative slots or depending on model-family
74/// names.
75#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct KvSlotCapacitySnapshot {
77    pub block_size: usize,
78    pub total_blocks: usize,
79    pub free_blocks: usize,
80}
81
82/// Token-validity mask for model-side greedy argmax.
83///
84/// `valid_token_mask[id] != 0` means token `id` may be selected. Tokens at or
85/// above `valid_token_mask.len()` are invalid. The fingerprint lets model
86/// backends cache an uploaded device mask without comparing the full vector on
87/// every decode step.
88#[derive(Clone)]
89pub struct TokenSelectionMask {
90    pub fingerprint: u64,
91    pub valid_token_mask: Arc<[i8]>,
92}
93
94impl TokenSelectionMask {
95    pub fn new(valid_token_mask: Vec<i8>) -> Self {
96        let fingerprint = Self::fingerprint(&valid_token_mask);
97        Self {
98            fingerprint,
99            valid_token_mask: Arc::from(valid_token_mask),
100        }
101    }
102
103    fn fingerprint(valid_token_mask: &[i8]) -> u64 {
104        let mut hasher = DefaultHasher::new();
105        valid_token_mask.hash(&mut hasher);
106        hasher.finish()
107    }
108
109    /// Change a small set of token-validity slots and refresh the cache key
110    /// once. `Arc::make_mut` keeps the common unshared request mask in place
111    /// and preserves safety when an in-flight backend policy still owns a
112    /// clone.
113    pub fn set_tokens_validity(&mut self, token_ids: &[u32], valid: bool) -> bool {
114        let value = i8::from(valid);
115        let slots = Arc::make_mut(&mut self.valid_token_mask);
116        let mut changed = false;
117        for &token_id in token_ids {
118            if let Some(slot) = slots.get_mut(token_id as usize) {
119                if *slot != value {
120                    *slot = value;
121                    changed = true;
122                }
123            }
124        }
125        if changed {
126            self.fingerprint = Self::fingerprint(slots);
127        }
128        changed
129    }
130
131    pub fn len(&self) -> usize {
132        self.valid_token_mask.len()
133    }
134
135    pub fn is_empty(&self) -> bool {
136        self.valid_token_mask.is_empty()
137    }
138}
139
140impl std::fmt::Debug for TokenSelectionMask {
141    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
142        let valid_count = self.valid_token_mask.iter().filter(|&&v| v != 0).count();
143        f.debug_struct("TokenSelectionMask")
144            .field("fingerprint", &self.fingerprint)
145            .field("len", &self.valid_token_mask.len())
146            .field("valid_count", &valid_count)
147            .finish()
148    }
149}
150
151#[cfg(test)]
152mod token_selection_mask_tests {
153    use super::TokenSelectionMask;
154
155    #[test]
156    fn response_completion_mask_is_copy_on_write_and_restores_fingerprint() {
157        let mut mask = TokenSelectionMask::new(vec![1, 1, 1]);
158        let original = mask.clone();
159        let original_fingerprint = mask.fingerprint;
160
161        assert!(mask.set_tokens_validity(&[1], false));
162        assert_eq!(mask.valid_token_mask.as_ref(), &[1, 0, 1]);
163        assert_eq!(original.valid_token_mask.as_ref(), &[1, 1, 1]);
164        assert_ne!(mask.fingerprint, original_fingerprint);
165
166        let masked_fingerprint = mask.fingerprint;
167        assert!(!mask.set_tokens_validity(&[1], false));
168        assert_eq!(mask.fingerprint, masked_fingerprint);
169
170        assert!(mask.set_tokens_validity(&[1], true));
171        assert_eq!(mask.fingerprint, original_fingerprint);
172    }
173}
174
175#[derive(Clone, Debug)]
176pub enum LogitsReturnPolicy {
177    FullLogits,
178    GreedyArgmax {
179        token_mask: Option<TokenSelectionMask>,
180        repetition_penalty: Option<GreedyRepetitionPenalty>,
181    },
182}
183
184impl Default for LogitsReturnPolicy {
185    fn default() -> Self {
186        Self::FullLogits
187    }
188}
189
190impl LogitsReturnPolicy {
191    pub fn requires_full_logits(&self) -> bool {
192        matches!(self, Self::FullLogits)
193    }
194}
195
196/// Typed product output returned by a plan-runtime execution wave.
197///
198/// A selected token is not a one-element logits vector. Keeping the variants
199/// distinct prevents the engine from inferring product semantics from tensor
200/// shape and lets full-logits requests retain every host-side sampler,
201/// grammar, and structured-output processor.
202#[derive(Debug, Clone, PartialEq)]
203pub enum ExecutorSamplingOutput {
204    FullLogits(Vec<f32>),
205    GreedyToken(TokenId),
206}
207
208impl ExecutorSamplingOutput {
209    pub fn full_logits(logits: Vec<f32>) -> Result<Self> {
210        if logits.is_empty() {
211            return Err(FerrumError::backend(
212                "plan-runtime sampling output requires non-empty logits",
213            ));
214        }
215        Ok(Self::FullLogits(logits))
216    }
217
218    pub const fn greedy_token(token: TokenId) -> Self {
219        Self::GreedyToken(token)
220    }
221
222    /// Validate that a device-selected token was explicitly authorized.
223    ///
224    /// Returning full logits for a greedy policy remains legal: heterogeneous
225    /// batches may deliberately fall back to host sampling. The inverse would
226    /// skip required product processors and therefore fails closed.
227    pub fn validate_for_policy(
228        &self,
229        policy: &LogitsReturnPolicy,
230        vocabulary_size: usize,
231    ) -> Result<()> {
232        match self {
233            Self::FullLogits(logits) if logits.len() != vocabulary_size => {
234                return Err(FerrumError::backend(format!(
235                    "plan runtime returned {} logits for vocabulary {vocabulary_size}",
236                    logits.len()
237                )));
238            }
239            Self::GreedyToken(_) if policy.requires_full_logits() => {
240                return Err(FerrumError::backend(
241                    "plan runtime returned a greedy token for a full-logits request",
242                ));
243            }
244            Self::GreedyToken(token)
245                if usize::try_from(token.get())
246                    .ok()
247                    .is_none_or(|token| token >= vocabulary_size) =>
248            {
249                return Err(FerrumError::backend(format!(
250                    "plan runtime returned token {} outside vocabulary {vocabulary_size}",
251                    token.get()
252                )));
253            }
254            _ => {}
255        }
256        Ok(())
257    }
258
259    pub fn into_full_logits(self) -> Result<Vec<f32>> {
260        match self {
261            Self::FullLogits(logits) => Ok(logits),
262            Self::GreedyToken(_) => Err(FerrumError::backend(
263                "plan-runtime prefill unexpectedly returned a selected token",
264            )),
265        }
266    }
267}
268
269#[cfg(test)]
270mod executor_sampling_output_tests {
271    use super::{ExecutorSamplingOutput, LogitsReturnPolicy};
272    use ferrum_types::TokenId;
273
274    #[test]
275    fn full_logits_require_exact_vocabulary_width() {
276        let output = ExecutorSamplingOutput::full_logits(vec![0.0; 4]).unwrap();
277        assert!(output
278            .validate_for_policy(&LogitsReturnPolicy::FullLogits, 4)
279            .is_ok());
280        assert!(output
281            .validate_for_policy(&LogitsReturnPolicy::FullLogits, 5)
282            .is_err());
283    }
284
285    #[test]
286    fn greedy_token_requires_greedy_policy_and_in_vocabulary_token() {
287        let allowed = LogitsReturnPolicy::GreedyArgmax {
288            token_mask: None,
289            repetition_penalty: None,
290        };
291        let output = ExecutorSamplingOutput::greedy_token(TokenId::new(3));
292        assert!(output.validate_for_policy(&allowed, 4).is_ok());
293        assert!(output.validate_for_policy(&allowed, 3).is_err());
294        assert!(output
295            .validate_for_policy(&LogitsReturnPolicy::FullLogits, 4)
296            .is_err());
297    }
298
299    #[test]
300    fn full_logits_are_a_legal_greedy_batch_fallback() {
301        let policy = LogitsReturnPolicy::GreedyArgmax {
302            token_mask: None,
303            repetition_penalty: None,
304        };
305        let output = ExecutorSamplingOutput::full_logits(vec![0.0; 4]).unwrap();
306        assert!(output.validate_for_policy(&policy, 4).is_ok());
307    }
308}
309
310/// Sparse repetition-penalty metadata for model-side greedy argmax.
311///
312/// The token list is request-local and de-duplicated. Applying the penalty
313/// before GPU argmax avoids downloading full `[batch, vocab]` logits for the
314/// common greedy chat path while preserving repeat avoidance.
315#[derive(Clone, Debug)]
316pub struct GreedyRepetitionPenalty {
317    penalty: f32,
318    token_ids: Arc<[u32]>,
319}
320
321impl GreedyRepetitionPenalty {
322    pub fn new(penalty: f32, mut token_ids: Vec<u32>) -> Self {
323        let mut seen = HashSet::with_capacity(token_ids.len());
324        token_ids.retain(|token| seen.insert(*token));
325        Self {
326            penalty,
327            token_ids: Arc::from(token_ids),
328        }
329    }
330
331    pub const fn penalty(&self) -> f32 {
332        self.penalty
333    }
334
335    pub fn token_ids(&self) -> &[u32] {
336        &self.token_ids
337    }
338
339    pub fn is_empty(&self) -> bool {
340        self.token_ids.is_empty() || self.penalty == 1.0
341    }
342}
343
344#[cfg(test)]
345mod greedy_repetition_penalty_tests {
346    use super::GreedyRepetitionPenalty;
347
348    #[test]
349    fn constructor_preserves_first_seen_order_and_removes_duplicates() {
350        let repetition = GreedyRepetitionPenalty::new(1.1, vec![7, 3, 7, 9, 3]);
351        assert_eq!(repetition.penalty(), 1.1);
352        assert_eq!(repetition.token_ids(), [7, 3, 9]);
353    }
354}
355
356/// Input for prefill phase (processing the initial prompt)
357#[derive(Debug, Clone)]
358pub struct PrefillInput {
359    /// Stable product request identity for plan-runtime resources.
360    pub request_id: Option<RequestId>,
361    /// Maximum sequence extent this request may reach, including the prompt.
362    /// Executors use this for fit validation without allocating future pages.
363    pub maximum_sequence_tokens: Option<usize>,
364    /// Exact scheduler-owned prompt chunk for this invocation.
365    ///
366    /// The input tensor still contains the full prompt so token identity and
367    /// global offsets remain stable. Plan runtimes execute only this range.
368    pub chunk: Option<PrefillChunk>,
369    /// Input token IDs [batch_size, sequence_length]
370    pub input_ids: TensorRef,
371    /// Attention mask [batch_size, sequence_length] (optional)
372    pub attention_mask: Option<TensorRef>,
373    /// Position IDs [batch_size, sequence_length] (optional, for RoPE)
374    pub position_ids: Option<TensorRef>,
375    /// Pre-allocated KV cache handle (optional, for paged attention)
376    pub kv_cache: Option<Arc<dyn KvCacheHandle>>,
377    /// Pre-allocated recurrent-state handle (optional, for state-space layers)
378    pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
379    /// Request metadata that can affect model execution.
380    pub metadata: HashMap<String, serde_json::Value>,
381}
382
383impl PrefillInput {
384    /// Create new prefill input
385    pub fn new(input_ids: TensorRef) -> Self {
386        Self {
387            request_id: None,
388            maximum_sequence_tokens: None,
389            chunk: None,
390            input_ids,
391            attention_mask: None,
392            position_ids: None,
393            kv_cache: None,
394            recurrent_state: None,
395            metadata: HashMap::new(),
396        }
397    }
398
399    /// Attach the typed request boundary consumed by plan runtimes.
400    pub fn with_request_context(
401        mut self,
402        request_id: RequestId,
403        maximum_sequence_tokens: usize,
404    ) -> Self {
405        self.request_id = Some(request_id);
406        self.maximum_sequence_tokens = Some(maximum_sequence_tokens);
407        self
408    }
409
410    /// Attach the exact scheduler-published prompt chunk.
411    pub fn with_chunk(mut self, chunk: PrefillChunk) -> Self {
412        self.chunk = Some(chunk);
413        self
414    }
415
416    /// Create prefill input with a pre-allocated KV cache handle.
417    pub fn with_kv_cache(mut self, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
418        self.kv_cache = Some(kv_cache);
419        self
420    }
421
422    /// Create prefill input with a pre-allocated recurrent-state handle.
423    pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
424        self.recurrent_state = Some(recurrent_state);
425        self
426    }
427
428    /// Attach request metadata.
429    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
430        self.metadata = metadata;
431        self
432    }
433
434    /// Add attention mask
435    pub fn with_attention_mask(mut self, mask: TensorRef) -> Self {
436        self.attention_mask = Some(mask);
437        self
438    }
439
440    /// Add position IDs
441    pub fn with_position_ids(mut self, positions: TensorRef) -> Self {
442        self.position_ids = Some(positions);
443        self
444    }
445
446    /// Get batch size
447    pub fn batch_size(&self) -> usize {
448        self.input_ids.shape()[0]
449    }
450
451    /// Get sequence length
452    pub fn sequence_length(&self) -> usize {
453        if self.input_ids.shape().len() >= 2 {
454            self.input_ids.shape()[1]
455        } else {
456            1
457        }
458    }
459}
460
461/// Exact, validated prompt progress assigned to one prefill invocation.
462#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
463pub struct PrefillChunk {
464    tokens_processed: usize,
465    tokens_to_process: usize,
466    total_prompt_tokens: usize,
467}
468
469#[cfg(test)]
470mod prefill_chunk_tests {
471    use super::PrefillChunk;
472
473    #[test]
474    fn validates_exact_progress_and_finality() {
475        let first = PrefillChunk::new(0, 3, 8).unwrap();
476        assert_eq!(first.range(), 0..3);
477        assert_eq!(first.end(), 3);
478        assert!(!first.is_final());
479
480        let final_chunk = PrefillChunk::new(3, 5, 8).unwrap();
481        assert_eq!(final_chunk.range(), 3..8);
482        assert!(final_chunk.is_final());
483    }
484
485    #[test]
486    fn rejects_empty_out_of_bounds_and_overflowing_progress() {
487        assert!(PrefillChunk::new(0, 0, 8).is_err());
488        assert!(PrefillChunk::new(0, 1, 0).is_err());
489        assert!(PrefillChunk::new(7, 2, 8).is_err());
490        assert!(PrefillChunk::new(usize::MAX, 1, usize::MAX).is_err());
491    }
492}
493
494impl PrefillChunk {
495    pub fn new(
496        tokens_processed: usize,
497        tokens_to_process: usize,
498        total_prompt_tokens: usize,
499    ) -> Result<Self> {
500        let end = tokens_processed
501            .checked_add(tokens_to_process)
502            .ok_or_else(|| {
503                ferrum_types::FerrumError::request_validation("prefill chunk overflows")
504            })?;
505        if tokens_to_process == 0 || total_prompt_tokens == 0 || end > total_prompt_tokens {
506            return Err(ferrum_types::FerrumError::request_validation(
507                "prefill chunk must be non-empty and within the full prompt",
508            ));
509        }
510        Ok(Self {
511            tokens_processed,
512            tokens_to_process,
513            total_prompt_tokens,
514        })
515    }
516
517    pub const fn tokens_processed(self) -> usize {
518        self.tokens_processed
519    }
520
521    pub const fn tokens_to_process(self) -> usize {
522        self.tokens_to_process
523    }
524
525    pub const fn total_prompt_tokens(self) -> usize {
526        self.total_prompt_tokens
527    }
528
529    pub fn range(self) -> Range<usize> {
530        self.tokens_processed..self.tokens_processed + self.tokens_to_process
531    }
532
533    pub const fn end(self) -> usize {
534        self.tokens_processed + self.tokens_to_process
535    }
536
537    pub const fn is_final(self) -> bool {
538        self.end() == self.total_prompt_tokens
539    }
540}
541
542/// Output from prefill phase
543#[derive(Debug, Clone)]
544pub struct PrefillOutput {
545    /// Logits for all positions [batch_size, sequence_length, vocab_size]
546    pub logits: TensorRef,
547    /// KV cache handle populated with prompt states
548    pub kv_cache: Arc<dyn KvCacheHandle>,
549    /// Recurrent-state handle populated with prompt state, when used.
550    pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
551    /// Hidden states at each layer (optional, for analysis)
552    pub hidden_states: Option<Vec<TensorRef>>,
553    /// Attention weights (optional, for analysis)
554    pub attention_weights: Option<Vec<TensorRef>>,
555}
556
557impl PrefillOutput {
558    /// Create new prefill output
559    pub fn new(logits: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
560        Self {
561            logits,
562            kv_cache,
563            recurrent_state: None,
564            hidden_states: None,
565            attention_weights: None,
566        }
567    }
568
569    /// Attach updated recurrent state to the prefill output.
570    pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
571        self.recurrent_state = Some(recurrent_state);
572        self
573    }
574
575    /// Get logits for last position (for next token generation)
576    pub fn last_token_logits(&self) -> Result<TensorRef> {
577        let shape = self.logits.shape();
578        if shape.len() != 3 {
579            return Err(ferrum_types::FerrumError::backend(
580                "Expected 3D logits tensor [batch, seq, vocab]",
581            ));
582        }
583
584        let seq_len = shape[1];
585        if seq_len == 0 {
586            return Err(ferrum_types::FerrumError::backend("Empty sequence"));
587        }
588
589        // Extract last position: [batch, seq-1:seq, vocab] -> [batch, vocab]
590        self.logits
591            .view(&[0, seq_len - 1, 0], &[shape[0], seq_len, shape[2]])
592    }
593}
594
595/// Input for decode phase (generating one token at a time)
596#[derive(Debug, Clone)]
597pub struct DecodeInput {
598    /// Stable product request identity for plan-runtime resources.
599    pub request_id: Option<RequestId>,
600    /// Input token ID for current step [batch_size, 1]
601    pub input_ids: TensorRef,
602    /// Existing KV cache from previous steps
603    pub kv_cache: Arc<dyn KvCacheHandle>,
604    /// Existing recurrent state from previous steps, when used.
605    pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
606    /// Position IDs for current step [batch_size, 1] (optional)
607    pub position_ids: Option<TensorRef>,
608    /// Request metadata that can affect model execution.
609    pub metadata: HashMap<String, serde_json::Value>,
610    /// How the model may return final-position logits for this request.
611    pub logits_policy: LogitsReturnPolicy,
612}
613
614impl DecodeInput {
615    /// Create new decode input
616    pub fn new(input_ids: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
617        Self {
618            request_id: None,
619            input_ids,
620            kv_cache,
621            recurrent_state: None,
622            position_ids: None,
623            metadata: HashMap::new(),
624            logits_policy: LogitsReturnPolicy::FullLogits,
625        }
626    }
627
628    /// Attach the product request identity to this decode step.
629    pub fn with_request_id(mut self, request_id: RequestId) -> Self {
630        self.request_id = Some(request_id);
631        self
632    }
633
634    /// Add position IDs
635    pub fn with_position_ids(mut self, positions: TensorRef) -> Self {
636        self.position_ids = Some(positions);
637        self
638    }
639
640    /// Attach recurrent state for state-space or hybrid layers.
641    pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
642        self.recurrent_state = Some(recurrent_state);
643        self
644    }
645
646    /// Attach request metadata.
647    pub fn with_metadata(mut self, metadata: HashMap<String, serde_json::Value>) -> Self {
648        self.metadata = metadata;
649        self
650    }
651
652    pub fn with_logits_policy(mut self, policy: LogitsReturnPolicy) -> Self {
653        self.logits_policy = policy;
654        self
655    }
656
657    /// Get batch size
658    pub fn batch_size(&self) -> usize {
659        self.input_ids.shape()[0]
660    }
661}
662
663/// One sequence's contribution to a unified mixed-batch forward.
664///
665/// A unified batch lets a single model forward pass process a mix of
666/// per-sequence work units: a prefill chunk (q_tokens.len() ≥ 1, possibly
667/// continuing from `pos_offset > 0` for chunked prefill) and a decode step
668/// (q_tokens.len() == 1, `pos_offset` = current cache length) coexist in
669/// the same call. The model layer concatenates all `q_tokens` into one
670/// [M_total, hidden] tensor and runs all GEMMs / norms once; only the
671/// attention kernel sees per-item segmentation.
672///
673/// This is the abstraction that enables vLLM-style chunked prefill where
674/// decode tokens for already-running sequences are produced in the same
675/// iter as a prefill chunk for a newly-arriving sequence.
676#[derive(Clone)]
677pub struct UnifiedBatchItem {
678    /// Identifier matching the sequence's KV cache (model-side keying).
679    pub seq_id: String,
680    /// Tokens to process this iter. For decode this is exactly 1 token;
681    /// for prefill (chunked or whole) this is the chunk's tokens.
682    pub q_tokens: Vec<u32>,
683    /// KV cache handle for this sequence.
684    pub kv_cache: Arc<dyn KvCacheHandle>,
685    /// Recurrent-state handle for this sequence, when used.
686    pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
687    /// Starting absolute position for the FIRST token in `q_tokens`.
688    /// 0 for a fresh prefill, `kv_len` for a decode step or a continuing
689    /// chunked-prefill slice.
690    pub pos_offset: usize,
691    /// True iff this item completes the request's prefill (or is a decode
692    /// item) — i.e. logits at the last token of `q_tokens` should be
693    /// returned for sampling. Intermediate prefill chunks set this false
694    /// to skip the lm_head + sampling path.
695    pub is_final_chunk: bool,
696    /// Request metadata that can affect model execution.
697    pub metadata: HashMap<String, serde_json::Value>,
698    /// How the model may return final-position logits for this item.
699    pub logits_policy: LogitsReturnPolicy,
700}
701
702impl std::fmt::Debug for UnifiedBatchItem {
703    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
704        f.debug_struct("UnifiedBatchItem")
705            .field("seq_id", &self.seq_id)
706            .field("q_len", &self.q_tokens.len())
707            .field("has_recurrent_state", &self.recurrent_state.is_some())
708            .field("pos_offset", &self.pos_offset)
709            .field("is_final_chunk", &self.is_final_chunk)
710            .finish()
711    }
712}
713
714/// A mixed-batch forward request: any combination of in-progress prefill
715/// chunks and decode steps. See [`UnifiedBatchItem`] for the per-item
716/// semantics. The producer (engine) groups all sequences active in this
717/// iter into a single batch; the consumer (model) runs one forward and
718/// returns per-item logits (only for items with `is_final_chunk = true`,
719/// in the order they appear in `items`).
720#[derive(Debug, Clone, Default)]
721pub struct UnifiedBatch {
722    pub items: Vec<UnifiedBatchItem>,
723}
724
725impl UnifiedBatch {
726    pub fn new() -> Self {
727        Self::default()
728    }
729
730    /// Total query tokens across all items — corresponds to the M dim of
731    /// the model's per-layer GEMMs in the unified forward.
732    pub fn total_q_tokens(&self) -> usize {
733        self.items.iter().map(|it| it.q_tokens.len()).sum()
734    }
735
736    /// Number of items that will produce a logits vector (decode items
737    /// always; prefill items only on their final chunk).
738    pub fn num_sampled_items(&self) -> usize {
739        self.items.iter().filter(|it| it.is_final_chunk).count()
740    }
741}
742
743/// Tensor-free decode input for an executor with plan-runtime resource
744/// authority.
745///
746/// The runtime already owns the request's physical cache and accepts exactly
747/// one token per decode frontier, so materializing a host tensor here adds no
748/// information. Legacy and modality executors continue to use [`DecodeInput`].
749#[derive(Debug, Clone)]
750pub struct PlanRuntimeDecodeInput {
751    pub request_id: RequestId,
752    pub input_token: TokenId,
753    pub kv_cache: Arc<dyn KvCacheHandle>,
754    pub logits_policy: LogitsReturnPolicy,
755}
756
757impl PlanRuntimeDecodeInput {
758    pub fn new(
759        request_id: RequestId,
760        input_token: TokenId,
761        kv_cache: Arc<dyn KvCacheHandle>,
762    ) -> Self {
763        Self {
764            request_id,
765            input_token,
766            kv_cache,
767            logits_policy: LogitsReturnPolicy::FullLogits,
768        }
769    }
770
771    pub fn with_logits_policy(mut self, logits_policy: LogitsReturnPolicy) -> Self {
772        self.logits_policy = logits_policy;
773        self
774    }
775}
776
777/// Tensor-free prefill input for an executor with plan-runtime resource
778/// authority.
779///
780/// The complete prompt remains immutable across chunk retries so scheduler
781/// offsets and executor sequence identity cannot drift. Physical KV and
782/// recurrent resources remain opaque executor-owned state.
783#[derive(Debug, Clone)]
784pub struct PlanRuntimePrefillInput {
785    pub request_id: RequestId,
786    pub input_tokens: Arc<[TokenId]>,
787    pub maximum_sequence_tokens: usize,
788    pub chunk: PrefillChunk,
789}
790
791impl PlanRuntimePrefillInput {
792    pub fn new(
793        request_id: RequestId,
794        input_tokens: impl Into<Arc<[TokenId]>>,
795        maximum_sequence_tokens: usize,
796        chunk: PrefillChunk,
797    ) -> Result<Self> {
798        let input_tokens = input_tokens.into();
799        if input_tokens.is_empty() {
800            return Err(FerrumError::request_validation(
801                "plan-runtime prefill requires at least one input token",
802            ));
803        }
804        if chunk.total_prompt_tokens() != input_tokens.len() {
805            return Err(FerrumError::request_validation(format!(
806                "plan-runtime prefill chunk declares {} prompt tokens for input length {}",
807                chunk.total_prompt_tokens(),
808                input_tokens.len()
809            )));
810        }
811        if maximum_sequence_tokens < input_tokens.len() {
812            return Err(FerrumError::request_validation(format!(
813                "plan-runtime sequence ceiling {maximum_sequence_tokens} does not cover prompt length {}",
814                input_tokens.len()
815            )));
816        }
817        Ok(Self {
818            request_id,
819            input_tokens,
820            maximum_sequence_tokens,
821            chunk,
822        })
823    }
824}
825
826/// Output from decode phase
827#[derive(Debug, Clone)]
828pub struct DecodeOutput {
829    /// Logits for next token [batch_size, vocab_size]
830    pub logits: TensorRef,
831    /// Updated KV cache with new token state
832    pub kv_cache: Arc<dyn KvCacheHandle>,
833    /// Updated recurrent state, when used.
834    pub recurrent_state: Option<Arc<dyn RecurrentStateHandle>>,
835    /// Hidden state for current token (optional)
836    pub hidden_state: Option<TensorRef>,
837    /// Attention weights for current token (optional)
838    pub attention_weights: Option<Vec<TensorRef>>,
839}
840
841impl DecodeOutput {
842    /// Create new decode output
843    pub fn new(logits: TensorRef, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
844        Self {
845            logits,
846            kv_cache,
847            recurrent_state: None,
848            hidden_state: None,
849            attention_weights: None,
850        }
851    }
852
853    /// Attach updated recurrent state to the decode output.
854    pub fn with_recurrent_state(mut self, recurrent_state: Arc<dyn RecurrentStateHandle>) -> Self {
855        self.recurrent_state = Some(recurrent_state);
856        self
857    }
858}
859
860/// Tensor-free decode output from a plan-runtime executor.
861#[derive(Debug, Clone)]
862pub struct PlanRuntimeDecodeOutput {
863    pub sampling_output: ExecutorSamplingOutput,
864    pub kv_cache: Arc<dyn KvCacheHandle>,
865}
866
867impl PlanRuntimeDecodeOutput {
868    pub fn new(sampling_output: ExecutorSamplingOutput, kv_cache: Arc<dyn KvCacheHandle>) -> Self {
869        Self {
870            sampling_output,
871            kv_cache,
872        }
873    }
874}
875
876/// Product state emitted by one completed plan-runtime prefill chunk.
877#[derive(Debug)]
878pub enum PlanRuntimePrefillProduct {
879    /// An intermediate chunk updates executor-owned state but is not sampleable.
880    Intermediate,
881    /// A final chunk returns the complete vocabulary row to the engine.
882    FinalLogits(Vec<f32>),
883}
884
885/// Exact executor-owned authority advanced by one prefill completion.
886#[derive(Debug)]
887pub struct PlanRuntimePrefillAuthority {
888    request_id: RequestId,
889    committed_tokens: usize,
890    kv_cache: Arc<dyn KvCacheHandle>,
891}
892
893impl PlanRuntimePrefillAuthority {
894    pub fn request_id(&self) -> &RequestId {
895        &self.request_id
896    }
897
898    pub const fn committed_tokens(&self) -> usize {
899        self.committed_tokens
900    }
901
902    pub fn kv_cache(&self) -> &Arc<dyn KvCacheHandle> {
903        &self.kv_cache
904    }
905
906    pub fn into_cache(self) -> Arc<dyn KvCacheHandle> {
907        self.kv_cache
908    }
909}
910
911/// Tensor-free prefill output bound to one request and committed KV extent.
912#[derive(Debug)]
913pub struct PlanRuntimePrefillOutput {
914    authority: PlanRuntimePrefillAuthority,
915    product: PlanRuntimePrefillProduct,
916}
917
918impl PlanRuntimePrefillOutput {
919    pub fn intermediate(
920        request_id: RequestId,
921        committed_tokens: usize,
922        kv_cache: Arc<dyn KvCacheHandle>,
923    ) -> Self {
924        Self {
925            authority: PlanRuntimePrefillAuthority {
926                request_id,
927                committed_tokens,
928                kv_cache,
929            },
930            product: PlanRuntimePrefillProduct::Intermediate,
931        }
932    }
933
934    pub fn final_logits(
935        request_id: RequestId,
936        committed_tokens: usize,
937        logits: Vec<f32>,
938        kv_cache: Arc<dyn KvCacheHandle>,
939    ) -> Result<Self> {
940        if logits.is_empty() {
941            return Err(FerrumError::backend(
942                "plan-runtime final prefill returned empty logits",
943            ));
944        }
945        Ok(Self {
946            authority: PlanRuntimePrefillAuthority {
947                request_id,
948                committed_tokens,
949                kv_cache,
950            },
951            product: PlanRuntimePrefillProduct::FinalLogits(logits),
952        })
953    }
954
955    pub fn request_id(&self) -> &RequestId {
956        self.authority.request_id()
957    }
958
959    pub const fn committed_tokens(&self) -> usize {
960        self.authority.committed_tokens()
961    }
962
963    pub fn product(&self) -> &PlanRuntimePrefillProduct {
964        &self.product
965    }
966
967    pub fn kv_cache(&self) -> &Arc<dyn KvCacheHandle> {
968        self.authority.kv_cache()
969    }
970
971    pub fn validate_for_completion(
972        &self,
973        expected_request_id: &RequestId,
974        completed_chunk: PrefillChunk,
975        vocabulary_size: usize,
976    ) -> Result<()> {
977        if self.request_id() != expected_request_id {
978            return Err(FerrumError::backend(format!(
979                "plan runtime returned prefill output for request {}, expected {expected_request_id}",
980                self.request_id()
981            )));
982        }
983        if self.committed_tokens() != completed_chunk.end() {
984            return Err(FerrumError::backend(format!(
985                "plan runtime returned prefill extent {}, expected {}",
986                self.committed_tokens(),
987                completed_chunk.end()
988            )));
989        }
990        if self.kv_cache().num_tokens() != self.committed_tokens() {
991            return Err(FerrumError::backend(format!(
992                "plan runtime prefill cache `{}` reports {} tokens for committed extent {}",
993                self.kv_cache().cache_id(),
994                self.kv_cache().num_tokens(),
995                self.committed_tokens()
996            )));
997        }
998        match (&self.product, completed_chunk.is_final()) {
999            (PlanRuntimePrefillProduct::Intermediate, false) => Ok(()),
1000            (PlanRuntimePrefillProduct::FinalLogits(logits), true)
1001                if logits.len() == vocabulary_size =>
1002            {
1003                Ok(())
1004            }
1005            (PlanRuntimePrefillProduct::FinalLogits(logits), true) => {
1006                Err(FerrumError::backend(format!(
1007                    "plan runtime returned {} final prefill logits for vocabulary {vocabulary_size}",
1008                    logits.len()
1009                )))
1010            }
1011            (PlanRuntimePrefillProduct::Intermediate, true) => Err(FerrumError::backend(
1012                "plan runtime returned an intermediate product for a final prefill chunk",
1013            )),
1014            (PlanRuntimePrefillProduct::FinalLogits(_), false) => Err(FerrumError::backend(
1015                "plan runtime returned final logits for an intermediate prefill chunk",
1016            )),
1017        }
1018    }
1019
1020    pub fn into_parts(self) -> (PlanRuntimePrefillAuthority, PlanRuntimePrefillProduct) {
1021        (self.authority, self.product)
1022    }
1023}
1024
1025/// Product-authoritative evidence for a successfully completed sequence.
1026///
1027/// Physical cache release is not completion evidence: cancellation, failure,
1028/// recompute, and successful generation all release the same cache authority.
1029/// The engine constructs this receipt only after it has finalized user-visible
1030/// token usage, allowing plan runtimes to reconcile terminal events without
1031/// inferring output counts from execution frames.
1032#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1033pub struct ExecutorSequenceCompletion {
1034    request_id: RequestId,
1035    cache_id: String,
1036    input_tokens: u64,
1037    output_tokens: u64,
1038}
1039
1040impl ExecutorSequenceCompletion {
1041    pub fn new(
1042        request_id: RequestId,
1043        cache_id: String,
1044        input_tokens: usize,
1045        output_tokens: usize,
1046    ) -> Result<Self> {
1047        if cache_id.is_empty() {
1048            return Err(FerrumError::request_validation(
1049                "executor sequence completion requires a cache identity",
1050            ));
1051        }
1052        let input_tokens = u64::try_from(input_tokens).map_err(|_| {
1053            FerrumError::request_validation("executor completion input token count exceeds u64")
1054        })?;
1055        let output_tokens = u64::try_from(output_tokens).map_err(|_| {
1056            FerrumError::request_validation("executor completion output token count exceeds u64")
1057        })?;
1058        Ok(Self {
1059            request_id,
1060            cache_id,
1061            input_tokens,
1062            output_tokens,
1063        })
1064    }
1065
1066    pub fn request_id(&self) -> &RequestId {
1067        &self.request_id
1068    }
1069
1070    pub fn cache_id(&self) -> &str {
1071        &self.cache_id
1072    }
1073
1074    pub const fn input_tokens(&self) -> u64 {
1075        self.input_tokens
1076    }
1077
1078    pub const fn output_tokens(&self) -> u64 {
1079        self.output_tokens
1080    }
1081}
1082
1083pub use ferrum_types::ExecutionResourceAuthority;
1084
1085/// Request-scoped authority selected for a capacity-pressure preemption.
1086///
1087/// The cache identity prevents the engine from releasing a newer sequence
1088/// incarnation after a stale scheduler decision. Implementations must retire
1089/// retained prefill and active decode authority through the same operation.
1090#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1091pub struct ExecutorExecutionCapacityPreemption {
1092    request_id: RequestId,
1093    cache_id: String,
1094}
1095
1096impl ExecutorExecutionCapacityPreemption {
1097    pub fn new(request_id: RequestId, cache_id: String) -> Result<Self> {
1098        if cache_id.is_empty() {
1099            return Err(FerrumError::request_validation(
1100                "execution-capacity preemption requires a cache identity",
1101            ));
1102        }
1103        Ok(Self {
1104            request_id,
1105            cache_id,
1106        })
1107    }
1108
1109    pub fn request_id(&self) -> &RequestId {
1110        &self.request_id
1111    }
1112
1113    pub fn cache_id(&self) -> &str {
1114        &self.cache_id
1115    }
1116}
1117
1118#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
1119#[serde(rename_all = "snake_case")]
1120pub enum ExecutorExecutionCapacityPreemptionAuthority {
1121    RetainedPrefill,
1122    ActiveSequence,
1123}
1124
1125/// Proof that the executor retired one exact request authority to a terminal
1126/// state. Source-generation advancement remains independently verified by the
1127/// engine before the scheduler may resume another frontier.
1128#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1129pub struct ExecutorExecutionCapacityPreemptionReceipt {
1130    request_id: RequestId,
1131    cache_id: String,
1132    authority: ExecutorExecutionCapacityPreemptionAuthority,
1133}
1134
1135impl ExecutorExecutionCapacityPreemptionReceipt {
1136    pub fn new(
1137        request_id: RequestId,
1138        cache_id: String,
1139        authority: ExecutorExecutionCapacityPreemptionAuthority,
1140    ) -> Self {
1141        Self {
1142            request_id,
1143            cache_id,
1144            authority,
1145        }
1146    }
1147
1148    pub fn request_id(&self) -> &RequestId {
1149        &self.request_id
1150    }
1151
1152    pub fn cache_id(&self) -> &str {
1153        &self.cache_id
1154    }
1155
1156    pub const fn authority(&self) -> ExecutorExecutionCapacityPreemptionAuthority {
1157        self.authority
1158    }
1159}
1160
1161/// Point-in-time memory evidence emitted by the shared plan runtime.
1162///
1163/// Static model allocations are separated from dynamic request resources so
1164/// product telemetry never reports model weights as KV or recurrent-state
1165/// usage. Process-wide claims are included because another live plan can
1166/// consume capacity visible to this runtime. Dynamic free bytes remain
1167/// reusable by this plan; quarantined and other claimed bytes do not.
1168#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1169pub struct PlanRuntimeResourceSnapshot {
1170    device_capacity_bytes: u64,
1171    usable_capacity_bytes: u64,
1172    process_claimed_bytes: u64,
1173    plan_claimed_bytes: u64,
1174    static_bytes: u64,
1175    dynamic_resident_bytes: u64,
1176    dynamic_free_bytes: u64,
1177    pending_growth_bytes: u64,
1178    quarantined_bytes: u64,
1179}
1180
1181impl PlanRuntimeResourceSnapshot {
1182    #[allow(clippy::too_many_arguments)]
1183    pub fn new(
1184        device_capacity_bytes: u64,
1185        usable_capacity_bytes: u64,
1186        process_claimed_bytes: u64,
1187        plan_claimed_bytes: u64,
1188        static_bytes: u64,
1189        dynamic_resident_bytes: u64,
1190        dynamic_free_bytes: u64,
1191        pending_growth_bytes: u64,
1192        quarantined_bytes: u64,
1193    ) -> Result<Self> {
1194        let snapshot = Self {
1195            device_capacity_bytes,
1196            usable_capacity_bytes,
1197            process_claimed_bytes,
1198            plan_claimed_bytes,
1199            static_bytes,
1200            dynamic_resident_bytes,
1201            dynamic_free_bytes,
1202            pending_growth_bytes,
1203            quarantined_bytes,
1204        };
1205        snapshot.validate()?;
1206        Ok(snapshot)
1207    }
1208
1209    /// Revalidates deserialized evidence before it crosses a trusted failure or
1210    /// profile boundary.
1211    pub fn validate(&self) -> Result<()> {
1212        if self.usable_capacity_bytes > self.device_capacity_bytes {
1213            return Err(ferrum_types::FerrumError::internal(format!(
1214                "plan runtime usable capacity {} exceeds device capacity {}",
1215                self.usable_capacity_bytes, self.device_capacity_bytes
1216            )));
1217        }
1218        if self.process_claimed_bytes > self.usable_capacity_bytes {
1219            return Err(ferrum_types::FerrumError::internal(format!(
1220                "plan runtime process claims {} exceed usable capacity {}",
1221                self.process_claimed_bytes, self.usable_capacity_bytes
1222            )));
1223        }
1224        if self.plan_claimed_bytes > self.process_claimed_bytes {
1225            return Err(ferrum_types::FerrumError::internal(format!(
1226                "plan runtime plan claims {} exceed process claims {}",
1227                self.plan_claimed_bytes, self.process_claimed_bytes
1228            )));
1229        }
1230        if self.dynamic_free_bytes > self.dynamic_resident_bytes {
1231            return Err(ferrum_types::FerrumError::internal(format!(
1232                "plan runtime dynamic free bytes {} exceed resident bytes {}",
1233                self.dynamic_free_bytes, self.dynamic_resident_bytes
1234            )));
1235        }
1236        let minimum_plan_claim = self
1237            .static_bytes
1238            .checked_add(self.dynamic_resident_bytes)
1239            .and_then(|bytes| bytes.checked_add(self.quarantined_bytes))
1240            .ok_or_else(|| {
1241                ferrum_types::FerrumError::internal(
1242                    "plan runtime static, resident, and quarantined bytes overflow u64",
1243                )
1244            })?;
1245        if minimum_plan_claim > self.plan_claimed_bytes {
1246            return Err(ferrum_types::FerrumError::internal(format!(
1247                "plan runtime accounted plan bytes {minimum_plan_claim} exceed plan claims {}",
1248                self.plan_claimed_bytes
1249            )));
1250        }
1251        Ok(())
1252    }
1253
1254    pub const fn device_capacity_bytes(&self) -> u64 {
1255        self.device_capacity_bytes
1256    }
1257
1258    pub const fn usable_capacity_bytes(&self) -> u64 {
1259        self.usable_capacity_bytes
1260    }
1261
1262    pub const fn process_claimed_bytes(&self) -> u64 {
1263        self.process_claimed_bytes
1264    }
1265
1266    pub const fn plan_claimed_bytes(&self) -> u64 {
1267        self.plan_claimed_bytes
1268    }
1269
1270    pub const fn static_bytes(&self) -> u64 {
1271        self.static_bytes
1272    }
1273
1274    pub const fn dynamic_resident_bytes(&self) -> u64 {
1275        self.dynamic_resident_bytes
1276    }
1277
1278    pub const fn dynamic_free_bytes(&self) -> u64 {
1279        self.dynamic_free_bytes
1280    }
1281
1282    pub const fn dynamic_used_bytes(&self) -> u64 {
1283        self.dynamic_resident_bytes - self.dynamic_free_bytes
1284    }
1285
1286    pub const fn pending_growth_bytes(&self) -> u64 {
1287        self.pending_growth_bytes
1288    }
1289
1290    pub const fn quarantined_bytes(&self) -> u64 {
1291        self.quarantined_bytes
1292    }
1293
1294    /// Capacity immediately reusable by this plan without reclaiming another
1295    /// plan: process-wide unclaimed bytes plus free extents already resident
1296    /// in this plan's dynamic pools.
1297    pub fn available_bytes(&self) -> Result<u64> {
1298        self.usable_capacity_bytes
1299            .checked_sub(self.process_claimed_bytes)
1300            .and_then(|bytes| bytes.checked_add(self.dynamic_free_bytes))
1301            .ok_or_else(|| {
1302                ferrum_types::FerrumError::internal(
1303                    "plan runtime available capacity calculation overflowed",
1304                )
1305            })
1306    }
1307
1308    pub fn used_bytes(&self) -> Result<u64> {
1309        self.available_bytes().and_then(|available| {
1310            self.usable_capacity_bytes
1311                .checked_sub(available)
1312                .ok_or_else(|| {
1313                    ferrum_types::FerrumError::internal(
1314                        "plan runtime available bytes exceed usable capacity",
1315                    )
1316                })
1317        })
1318    }
1319}
1320
1321#[cfg(test)]
1322mod plan_runtime_resource_snapshot_tests {
1323    use super::PlanRuntimeResourceSnapshot;
1324
1325    #[test]
1326    fn separates_static_and_dynamic_usage() {
1327        let snapshot =
1328            PlanRuntimeResourceSnapshot::new(1_000, 900, 710, 710, 400, 300, 200, 20, 10).unwrap();
1329
1330        assert_eq!(snapshot.available_bytes().unwrap(), 390);
1331        assert_eq!(snapshot.used_bytes().unwrap(), 510);
1332        assert_eq!(snapshot.dynamic_resident_bytes(), 300);
1333        assert_eq!(snapshot.dynamic_used_bytes(), 100);
1334        assert_eq!(snapshot.dynamic_free_bytes(), 200);
1335        assert_eq!(snapshot.pending_growth_bytes(), 20);
1336        assert_eq!(snapshot.quarantined_bytes(), 10);
1337    }
1338
1339    #[test]
1340    fn rejects_incoherent_capacity_evidence() {
1341        assert!(PlanRuntimeResourceSnapshot::new(1_000, 1_001, 0, 0, 0, 0, 0, 0, 0).is_err());
1342        assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 901, 0, 0, 0, 0, 0, 0).is_err());
1343        assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 500, 501, 0, 0, 0, 0, 0).is_err());
1344        assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 100, 100, 0, 100, 101, 0, 0).is_err());
1345        assert!(PlanRuntimeResourceSnapshot::new(1_000, 900, 500, 500, 400, 100, 0, 0, 1).is_err());
1346    }
1347}
1348
1349/// Typed origin for an executor-owned request lifecycle.
1350#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1351pub enum ExecutorRequestOrigin {
1352    Product,
1353    Startup,
1354    Diagnostic,
1355}
1356
1357impl ExecutorRequestOrigin {
1358    pub const fn namespace(self) -> &'static str {
1359        match self {
1360            Self::Product => "product",
1361            Self::Startup => "startup",
1362            Self::Diagnostic => "diagnostic",
1363        }
1364    }
1365
1366    pub fn from_namespaced_request_identity(identity: &str) -> Option<Self> {
1367        let suffix = identity.strip_prefix("request.")?;
1368        let (namespace, request_id) = suffix.split_once('.')?;
1369        if request_id.is_empty() {
1370            return None;
1371        }
1372        match namespace {
1373            "product" => Some(Self::Product),
1374            "startup" => Some(Self::Startup),
1375            "diagnostic" => Some(Self::Diagnostic),
1376            _ => None,
1377        }
1378    }
1379}
1380
1381/// Borrowed, already-tokenized input used to probe plan-runtime prefill
1382/// admission before the request can enter a device submission batch.
1383///
1384/// This carries semantic token identity rather than an aggregate token count:
1385/// vNext derives the exact resource work shape and its fingerprint from this
1386/// boundary. The request remains owned by the scheduler while the executor
1387/// retains any admitted authority internally until [`ModelExecutor::prefill`]
1388/// consumes it or cancellation releases it.
1389#[derive(Debug, Clone, Copy)]
1390pub struct ExecutorPrefillAdmission<'a> {
1391    pub request_id: &'a RequestId,
1392    pub input_tokens: &'a [TokenId],
1393    pub maximum_sequence_tokens: usize,
1394    /// Prompt tokens reported in the final product usage.
1395    pub product_prompt_tokens: usize,
1396    /// Already-committed output tokens replayed as part of a recompute input.
1397    pub replayed_output_tokens: usize,
1398    pub request_origin: ExecutorRequestOrigin,
1399}
1400
1401impl<'a> ExecutorPrefillAdmission<'a> {
1402    pub const fn for_startup(
1403        request_id: &'a RequestId,
1404        input_tokens: &'a [TokenId],
1405        maximum_sequence_tokens: usize,
1406    ) -> Self {
1407        Self {
1408            request_id,
1409            input_tokens,
1410            maximum_sequence_tokens,
1411            product_prompt_tokens: input_tokens.len(),
1412            replayed_output_tokens: 0,
1413            request_origin: ExecutorRequestOrigin::Startup,
1414        }
1415    }
1416
1417    pub const fn for_diagnostic(
1418        request_id: &'a RequestId,
1419        input_tokens: &'a [TokenId],
1420        maximum_sequence_tokens: usize,
1421    ) -> Self {
1422        Self {
1423            request_id,
1424            input_tokens,
1425            maximum_sequence_tokens,
1426            product_prompt_tokens: input_tokens.len(),
1427            replayed_output_tokens: 0,
1428            request_origin: ExecutorRequestOrigin::Diagnostic,
1429        }
1430    }
1431
1432    /// Construct admission for a product request whose execution context may
1433    /// include output tokens replayed after preemption.
1434    pub fn for_product_request(
1435        request_id: &'a RequestId,
1436        input_tokens: &'a [TokenId],
1437        maximum_sequence_tokens: usize,
1438        product_prompt_tokens: usize,
1439        replayed_output_tokens: usize,
1440    ) -> Result<Self> {
1441        let admission = Self {
1442            request_id,
1443            input_tokens,
1444            maximum_sequence_tokens,
1445            product_prompt_tokens,
1446            replayed_output_tokens,
1447            request_origin: ExecutorRequestOrigin::Product,
1448        };
1449        admission.validate()?;
1450        Ok(admission)
1451    }
1452
1453    pub fn validate(&self) -> Result<()> {
1454        if self.input_tokens.is_empty() {
1455            return Err(FerrumError::request_validation(
1456                "executor prefill admission requires at least one execution-context token",
1457            ));
1458        }
1459        if self.product_prompt_tokens == 0 {
1460            return Err(FerrumError::request_validation(
1461                "executor prefill admission requires at least one product prompt token",
1462            ));
1463        }
1464        let execution_context_tokens = self
1465            .product_prompt_tokens
1466            .checked_add(self.replayed_output_tokens)
1467            .ok_or_else(|| {
1468                FerrumError::request_validation(
1469                    "executor prefill product token accounting exceeds usize",
1470                )
1471            })?;
1472        if execution_context_tokens != self.input_tokens.len() {
1473            return Err(FerrumError::request_validation(format!(
1474                "executor prefill execution context has {} tokens but product accounting declares {} prompt + {} replayed output",
1475                self.input_tokens.len(),
1476                self.product_prompt_tokens,
1477                self.replayed_output_tokens
1478            )));
1479        }
1480        if self.maximum_sequence_tokens < execution_context_tokens {
1481            return Err(FerrumError::request_validation(format!(
1482                "executor prefill sequence ceiling {} does not cover execution context {execution_context_tokens}",
1483                self.maximum_sequence_tokens
1484            )));
1485        }
1486        Ok(())
1487    }
1488}
1489
1490#[cfg(test)]
1491mod executor_prefill_admission_tests {
1492    use super::{ExecutorPrefillAdmission, ExecutorRequestOrigin};
1493    use ferrum_types::{RequestId, TokenId};
1494
1495    #[test]
1496    fn product_accounting_distinguishes_replayed_output_from_prompt() {
1497        let request_id = RequestId::new();
1498        let tokens = [1, 2, 3, 4, 5]
1499            .into_iter()
1500            .map(TokenId::new)
1501            .collect::<Vec<_>>();
1502
1503        let admission =
1504            ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 8, 3, 2)
1505                .expect("recompute accounting must be accepted");
1506
1507        assert_eq!(admission.product_prompt_tokens, 3);
1508        assert_eq!(admission.replayed_output_tokens, 2);
1509        assert_eq!(admission.request_origin, ExecutorRequestOrigin::Product);
1510        assert_eq!(
1511            ExecutorPrefillAdmission::for_startup(&request_id, &tokens, 8).request_origin,
1512            ExecutorRequestOrigin::Startup
1513        );
1514        assert_eq!(
1515            ExecutorPrefillAdmission::for_diagnostic(&request_id, &tokens, 8).request_origin,
1516            ExecutorRequestOrigin::Diagnostic
1517        );
1518        assert_eq!(ExecutorRequestOrigin::Product.namespace(), "product");
1519        assert_eq!(ExecutorRequestOrigin::Startup.namespace(), "startup");
1520        assert_eq!(ExecutorRequestOrigin::Diagnostic.namespace(), "diagnostic");
1521        assert_eq!(
1522            ExecutorRequestOrigin::from_namespaced_request_identity("request.product.123"),
1523            Some(ExecutorRequestOrigin::Product)
1524        );
1525        assert_eq!(
1526            ExecutorRequestOrigin::from_namespaced_request_identity("request.startup.123"),
1527            Some(ExecutorRequestOrigin::Startup)
1528        );
1529        assert_eq!(
1530            ExecutorRequestOrigin::from_namespaced_request_identity("request.diagnostic.123"),
1531            Some(ExecutorRequestOrigin::Diagnostic)
1532        );
1533        assert_eq!(
1534            ExecutorRequestOrigin::from_namespaced_request_identity("request.product."),
1535            None
1536        );
1537        assert_eq!(
1538            ExecutorRequestOrigin::from_namespaced_request_identity("request/external"),
1539            None
1540        );
1541    }
1542
1543    #[test]
1544    fn product_accounting_rejects_context_drift_and_short_ceiling() {
1545        let request_id = RequestId::new();
1546        let tokens = [1, 2, 3].into_iter().map(TokenId::new).collect::<Vec<_>>();
1547
1548        assert!(
1549            ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 3, 2, 0).is_err()
1550        );
1551        assert!(
1552            ExecutorPrefillAdmission::for_product_request(&request_id, &tokens, 2, 2, 1).is_err()
1553        );
1554    }
1555}
1556
1557/// Scheduler-visible proof that an executor retained request and sequence
1558/// authority for future scheduler-owned prefill chunks.
1559#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1560pub struct ExecutorPrefillAdmissionReceipt {
1561    pub request_id: RequestId,
1562}
1563
1564/// Stable scheduler-facing projection of one plan-runtime capacity domain.
1565#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1566pub struct ExecutorAdmissionEpochs {
1567    pub coordinator_id: NonZeroU64,
1568    pub release_epoch: u64,
1569    pub capacity_epoch: u64,
1570}
1571
1572impl ExecutorAdmissionEpochs {
1573    pub const fn new(coordinator_id: NonZeroU64, release_epoch: u64, capacity_epoch: u64) -> Self {
1574        Self {
1575            coordinator_id,
1576            release_epoch,
1577            capacity_epoch,
1578        }
1579    }
1580
1581    pub fn from_capacity(epochs: crate::vnext::CapacityEpochs) -> Self {
1582        Self::new(
1583            NonZeroU64::new(epochs.coordinator_id().get())
1584                .expect("core-issued admission coordinator ids are non-zero"),
1585            epochs.release_epoch(),
1586            epochs.capacity_epoch(),
1587        )
1588    }
1589}
1590
1591type ExecutorCapacityWaitFuture =
1592    Pin<Box<dyn Future<Output = Result<ExecutorAdmissionEpochs>> + Send + 'static>>;
1593
1594/// Type-erased, single-use registration for one plan-runtime capacity wait.
1595///
1596/// Registration is created synchronously so the executor can subscribe before
1597/// the engine releases its iteration lock. Awaiting it never grants resources;
1598/// it only returns fresh epochs that permit another authoritative admission
1599/// probe.
1600#[must_use = "capacity wait registrations must be awaited or explicitly dropped"]
1601pub struct ExecutorCapacityWaitRegistration {
1602    future: ExecutorCapacityWaitFuture,
1603}
1604
1605impl ExecutorCapacityWaitRegistration {
1606    pub fn new<F>(future: F) -> Self
1607    where
1608        F: Future<Output = Result<ExecutorAdmissionEpochs>> + Send + 'static,
1609    {
1610        Self {
1611            future: Box::pin(future),
1612        }
1613    }
1614
1615    pub async fn wait_for_change(self) -> Result<ExecutorAdmissionEpochs> {
1616        self.future.await
1617    }
1618}
1619
1620/// Pre-submit runtime stage that could not acquire its exact dynamic capacity.
1621#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1622#[serde(rename_all = "snake_case")]
1623pub enum ExecutorExecutionCapacityStage {
1624    SequenceExtension,
1625    StepAdmission,
1626    SubmissionWave,
1627}
1628
1629/// One exact physical pool mutation completed while an executor was trying to
1630/// make an unsubmitted frontier runnable.
1631#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1632pub struct ExecutorExecutionMaintenanceMutation {
1633    pool_id: crate::vnext::DynamicBackingPoolId,
1634    domain_id: crate::vnext::CapacityDomainId,
1635    chunk: crate::vnext::BackingChunkIdentity,
1636    chunk_bytes: u64,
1637    published_capacity_bytes: u64,
1638    capacity_epoch: u64,
1639}
1640
1641impl ExecutorExecutionMaintenanceMutation {
1642    pub fn pool_id(&self) -> &crate::vnext::DynamicBackingPoolId {
1643        &self.pool_id
1644    }
1645
1646    pub const fn domain_id(&self) -> crate::vnext::CapacityDomainId {
1647        self.domain_id
1648    }
1649
1650    pub fn chunk(&self) -> &crate::vnext::BackingChunkIdentity {
1651        &self.chunk
1652    }
1653
1654    pub const fn chunk_bytes(&self) -> u64 {
1655        self.chunk_bytes
1656    }
1657
1658    pub const fn published_capacity_bytes(&self) -> u64 {
1659        self.published_capacity_bytes
1660    }
1661
1662    pub const fn capacity_epoch(&self) -> u64 {
1663        self.capacity_epoch
1664    }
1665}
1666
1667/// Typed proof that a bounded executor call committed relevant backing growth
1668/// before yielding the same logical frontier back to the scheduler.
1669///
1670/// This is deliberately stronger than observing a global capacity epoch. Every
1671/// mutation is reconstructed from a real growth receipt and bound back to the
1672/// pool's exact capacity domain. The scheduler may grant one fairness-bounded
1673/// retry only when this proof is present.
1674#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1675pub struct ExecutorExecutionMaintenanceProgress {
1676    attempts: u32,
1677    coordinator_id: NonZeroU64,
1678    mutations: Vec<ExecutorExecutionMaintenanceMutation>,
1679    latest_capacity_epoch: u64,
1680}
1681
1682impl ExecutorExecutionMaintenanceProgress {
1683    pub fn from_growth_receipts(
1684        attempts: u32,
1685        observed: ExecutorAdmissionEpochs,
1686        receipts: &[crate::vnext::DynamicPoolGrowthBatchReceipt],
1687        pools: &[crate::vnext::DynamicPoolStatus],
1688    ) -> Result<Self> {
1689        if attempts == 0 || receipts.is_empty() || receipts.len() > attempts as usize {
1690            return Err(FerrumError::internal(
1691                "execution maintenance retry requires bounded, non-empty growth receipts",
1692            ));
1693        }
1694
1695        let mut mutations = Vec::new();
1696        let mut previous_capacity_epoch = None;
1697        for receipt in receipts {
1698            if receipt.coordinator_id().get() != observed.coordinator_id.get() {
1699                return Err(FerrumError::internal(
1700                    "execution maintenance receipt belongs to another capacity coordinator",
1701                ));
1702            }
1703            if receipt.growths().is_empty()
1704                || previous_capacity_epoch
1705                    .is_some_and(|previous| receipt.capacity_epoch() <= previous)
1706            {
1707                return Err(FerrumError::internal(
1708                    "execution maintenance receipts contain no new ordered capacity mutation",
1709                ));
1710            }
1711            previous_capacity_epoch = Some(receipt.capacity_epoch());
1712
1713            for growth in receipt.growths() {
1714                let pool = pools
1715                    .iter()
1716                    .find(|pool| pool.pool_id() == growth.pool_id())
1717                    .ok_or_else(|| {
1718                        FerrumError::internal(
1719                            "execution maintenance receipt references an unknown dynamic pool",
1720                        )
1721                    })?;
1722                if growth.chunk().pool_id() != growth.pool_id()
1723                    || growth.chunk_bytes() == 0
1724                    || growth.published_capacity_bytes() == 0
1725                    || growth.capacity_epoch() != receipt.capacity_epoch()
1726                {
1727                    return Err(FerrumError::internal(
1728                        "execution maintenance receipt contains an invalid pool mutation",
1729                    ));
1730                }
1731                if mutations
1732                    .iter()
1733                    .any(|mutation: &ExecutorExecutionMaintenanceMutation| {
1734                        mutation.pool_id() == growth.pool_id() && mutation.chunk() == growth.chunk()
1735                    })
1736                {
1737                    return Err(FerrumError::internal(
1738                        "execution maintenance receipts repeat one physical pool mutation",
1739                    ));
1740                }
1741                mutations.push(ExecutorExecutionMaintenanceMutation {
1742                    pool_id: growth.pool_id().clone(),
1743                    domain_id: pool.domain_id(),
1744                    chunk: growth.chunk().clone(),
1745                    chunk_bytes: growth.chunk_bytes(),
1746                    published_capacity_bytes: growth.published_capacity_bytes(),
1747                    capacity_epoch: growth.capacity_epoch(),
1748                });
1749            }
1750        }
1751
1752        let latest_capacity_epoch = previous_capacity_epoch.expect("receipts are non-empty");
1753        if latest_capacity_epoch > observed.capacity_epoch {
1754            return Err(FerrumError::internal(
1755                "execution maintenance receipt is newer than the exported capacity observation",
1756            ));
1757        }
1758        mutations.sort_by(|left, right| {
1759            (
1760                left.capacity_epoch,
1761                left.pool_id.as_str(),
1762                left.chunk.ordinal(),
1763                left.chunk.generation(),
1764            )
1765                .cmp(&(
1766                    right.capacity_epoch,
1767                    right.pool_id.as_str(),
1768                    right.chunk.ordinal(),
1769                    right.chunk.generation(),
1770                ))
1771        });
1772        Ok(Self {
1773            attempts,
1774            coordinator_id: observed.coordinator_id,
1775            mutations,
1776            latest_capacity_epoch,
1777        })
1778    }
1779
1780    pub const fn attempts(&self) -> u32 {
1781        self.attempts
1782    }
1783
1784    pub const fn coordinator_id(&self) -> NonZeroU64 {
1785        self.coordinator_id
1786    }
1787
1788    pub fn mutations(&self) -> &[ExecutorExecutionMaintenanceMutation] {
1789        &self.mutations
1790    }
1791
1792    pub const fn latest_capacity_epoch(&self) -> u64 {
1793        self.latest_capacity_epoch
1794    }
1795}
1796
1797/// Scheduler-consumable proof binding physical maintenance to the exact
1798/// logical frontiers whose unsubmitted execution attempt observed it.
1799///
1800/// Keeping the request scope inside the proof prevents a caller from attaching
1801/// one sequence's growth receipt to an unrelated decode cohort.
1802#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1803pub struct ExecutorExecutionMaintenanceRetry {
1804    affected_request_ids: Vec<RequestId>,
1805    progress: ExecutorExecutionMaintenanceProgress,
1806}
1807
1808impl ExecutorExecutionMaintenanceRetry {
1809    fn new(
1810        affected_request_ids: Vec<RequestId>,
1811        progress: ExecutorExecutionMaintenanceProgress,
1812    ) -> Result<Self> {
1813        let unique = affected_request_ids.iter().collect::<HashSet<_>>();
1814        if affected_request_ids.is_empty() || unique.len() != affected_request_ids.len() {
1815            return Err(FerrumError::internal(
1816                "execution maintenance retry requires unique affected requests",
1817            ));
1818        }
1819        if progress.mutations().is_empty() {
1820            return Err(FerrumError::internal(
1821                "execution maintenance retry requires physical mutations",
1822            ));
1823        }
1824        Ok(Self {
1825            affected_request_ids,
1826            progress,
1827        })
1828    }
1829
1830    pub fn affected_request_ids(&self) -> &[RequestId] {
1831        &self.affected_request_ids
1832    }
1833
1834    pub const fn progress(&self) -> &ExecutorExecutionMaintenanceProgress {
1835        &self.progress
1836    }
1837}
1838
1839/// Scheduler-visible proof that an execution attempt was not submitted and
1840/// must not be retried until one of its exact capacity sources changes.
1841///
1842/// This value owns no resource authority. The executor retains the committed
1843/// request/sequence authority and has already retired any unsubmitted step or
1844/// submission-wave authority before returning it.
1845#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
1846#[serde(rename_all = "snake_case")]
1847pub enum ExecutorExecutionCapacityEvidenceOwner {
1848    Logical,
1849    Backing,
1850}
1851
1852#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1853#[serde(tag = "kind", rename_all = "snake_case")]
1854enum ExecutorExecutionCapacityEvidenceKind {
1855    Logical {
1856        shortfalls: Vec<crate::vnext::CapacityShortfall>,
1857        #[serde(skip_serializing_if = "Option::is_none")]
1858        pressure: Option<crate::vnext::DynamicBackingPressure>,
1859    },
1860    BackingDeferred {
1861        blockers: Vec<crate::vnext::DynamicBackingBlocker>,
1862        #[serde(skip_serializing_if = "Option::is_none")]
1863        pressure: Option<crate::vnext::DynamicBackingPressure>,
1864    },
1865    BackingPressure {
1866        pressure: crate::vnext::DynamicBackingPressure,
1867    },
1868}
1869
1870/// Exactly one typed owner for an execution-capacity deferral.
1871///
1872/// Its fields are private so callers cannot construct an empty or ambiguous
1873/// logical/backing evidence combination.
1874#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
1875pub struct ExecutorExecutionCapacityEvidence {
1876    owner: ExecutorExecutionCapacityEvidenceOwner,
1877    #[serde(flatten)]
1878    kind: ExecutorExecutionCapacityEvidenceKind,
1879    #[serde(skip_serializing_if = "Option::is_none")]
1880    maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
1881}
1882
1883impl ExecutorExecutionCapacityEvidence {
1884    fn logical(shortfalls: Vec<crate::vnext::CapacityShortfall>) -> Result<Self> {
1885        Self::logical_with_pressure(shortfalls, None)
1886    }
1887
1888    fn logical_with_pressure(
1889        shortfalls: Vec<crate::vnext::CapacityShortfall>,
1890        pressure: Option<crate::vnext::DynamicBackingPressure>,
1891    ) -> Result<Self> {
1892        if shortfalls.is_empty() {
1893            return Err(FerrumError::internal(
1894                "logical execution deferral requires at least one shortfall",
1895            ));
1896        }
1897        Ok(Self {
1898            owner: ExecutorExecutionCapacityEvidenceOwner::Logical,
1899            kind: ExecutorExecutionCapacityEvidenceKind::Logical {
1900                shortfalls,
1901                pressure,
1902            },
1903            maintenance_boundary: None,
1904        })
1905    }
1906
1907    fn backing_deferred(blockers: Vec<crate::vnext::DynamicBackingBlocker>) -> Result<Self> {
1908        Self::backing_deferred_with_pressure(blockers, None)
1909    }
1910
1911    fn backing_deferred_with_pressure(
1912        blockers: Vec<crate::vnext::DynamicBackingBlocker>,
1913        pressure: Option<crate::vnext::DynamicBackingPressure>,
1914    ) -> Result<Self> {
1915        if blockers.is_empty() {
1916            return Err(FerrumError::internal(
1917                "physical execution deferral requires at least one backing blocker",
1918            ));
1919        }
1920        Ok(Self {
1921            owner: ExecutorExecutionCapacityEvidenceOwner::Backing,
1922            kind: ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, pressure },
1923            maintenance_boundary: None,
1924        })
1925    }
1926
1927    fn direct_backing_pressure(pressure: crate::vnext::DynamicBackingPressure) -> Self {
1928        Self {
1929            owner: ExecutorExecutionCapacityEvidenceOwner::Backing,
1930            kind: ExecutorExecutionCapacityEvidenceKind::BackingPressure { pressure },
1931            maintenance_boundary: None,
1932        }
1933    }
1934
1935    pub const fn owner(&self) -> ExecutorExecutionCapacityEvidenceOwner {
1936        self.owner
1937    }
1938
1939    pub fn shortfalls(&self) -> &[crate::vnext::CapacityShortfall] {
1940        match &self.kind {
1941            ExecutorExecutionCapacityEvidenceKind::Logical { shortfalls, .. } => shortfalls,
1942            ExecutorExecutionCapacityEvidenceKind::BackingDeferred { .. }
1943            | ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => &[],
1944        }
1945    }
1946
1947    pub fn backing_blockers(&self) -> &[crate::vnext::DynamicBackingBlocker] {
1948        match &self.kind {
1949            ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, .. } => blockers,
1950            ExecutorExecutionCapacityEvidenceKind::Logical { .. }
1951            | ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => &[],
1952        }
1953    }
1954
1955    pub const fn backing_pressure(&self) -> Option<&crate::vnext::DynamicBackingPressure> {
1956        match &self.kind {
1957            ExecutorExecutionCapacityEvidenceKind::Logical { pressure, .. }
1958            | ExecutorExecutionCapacityEvidenceKind::BackingDeferred { pressure, .. } => {
1959                pressure.as_ref()
1960            }
1961            ExecutorExecutionCapacityEvidenceKind::BackingPressure { pressure } => Some(pressure),
1962        }
1963    }
1964
1965    pub const fn maintenance_boundary(
1966        &self,
1967    ) -> Option<&crate::vnext::DynamicPoolMaintenanceBoundaryReceipt> {
1968        self.maintenance_boundary.as_ref()
1969    }
1970
1971    fn with_maintenance_boundary(
1972        mut self,
1973        boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
1974    ) -> Result<Self> {
1975        match (self.backing_pressure(), boundary.as_ref()) {
1976            (
1977                Some(crate::vnext::DynamicBackingPressure::DeviceCapacity(pressure)),
1978                Some(boundary),
1979            ) if pressure == boundary.pressure() && !boundary.reclaim_sufficient() => {}
1980            (Some(crate::vnext::DynamicBackingPressure::PoolResident(_)), None) | (None, None) => {}
1981            (Some(crate::vnext::DynamicBackingPressure::DeviceCapacity(_)), None) => {
1982                return Err(FerrumError::internal(
1983                    "device-capacity execution maintenance lost its boundary receipt",
1984                ));
1985            }
1986            _ => {
1987                return Err(FerrumError::internal(
1988                    "execution maintenance boundary differs from its blocked pressure",
1989                ));
1990            }
1991        }
1992        self.maintenance_boundary = boundary;
1993        Ok(self)
1994    }
1995
1996    fn has_relevant_mutation(&self, mutation: &ExecutorExecutionMaintenanceMutation) -> bool {
1997        let logical_matches = |shortfalls: &[crate::vnext::CapacityShortfall]| {
1998            shortfalls.iter().any(|shortfall| {
1999                shortfall.kind() == crate::vnext::CapacityShortfallKind::BackingGrowthRequired
2000                    && shortfall.domain() == Some(mutation.domain_id())
2001            })
2002        };
2003        let backing_matches = |blockers: &[crate::vnext::DynamicBackingBlocker]| {
2004            blockers.iter().any(|blocker| {
2005                blocker.pool_id() == mutation.pool_id()
2006                    && blocker.domain_id() == mutation.domain_id()
2007            })
2008        };
2009        match &self.kind {
2010            ExecutorExecutionCapacityEvidenceKind::Logical { shortfalls, .. } => {
2011                logical_matches(shortfalls)
2012            }
2013            ExecutorExecutionCapacityEvidenceKind::BackingDeferred { blockers, .. } => {
2014                backing_matches(blockers)
2015            }
2016            ExecutorExecutionCapacityEvidenceKind::BackingPressure { .. } => false,
2017        }
2018    }
2019}
2020
2021#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2022pub struct ExecutorExecutionCapacityDeferral {
2023    observed: ExecutorAdmissionEpochs,
2024    wait_condition: crate::vnext::CapacityWaitCondition,
2025    stage: ExecutorExecutionCapacityStage,
2026    evidence: ExecutorExecutionCapacityEvidence,
2027    #[serde(skip_serializing_if = "Option::is_none")]
2028    maintenance_retry: Option<ExecutorExecutionMaintenanceRetry>,
2029}
2030
2031impl ExecutorExecutionCapacityDeferral {
2032    fn with_evidence(
2033        observed: ExecutorAdmissionEpochs,
2034        wait_condition: crate::vnext::CapacityWaitCondition,
2035        stage: ExecutorExecutionCapacityStage,
2036        evidence: ExecutorExecutionCapacityEvidence,
2037    ) -> Result<Self> {
2038        if wait_condition.coordinator_id().get() != observed.coordinator_id.get() {
2039            return Err(ferrum_types::FerrumError::request_validation(
2040                "executor execution deferral belongs to a different capacity coordinator",
2041            ));
2042        }
2043        Ok(Self {
2044            observed,
2045            wait_condition,
2046            stage,
2047            evidence,
2048            maintenance_retry: None,
2049        })
2050    }
2051
2052    /// Construct a physical deferral for an executor that observes device or
2053    /// pool pressure directly rather than through Ferrum's backing allocator.
2054    pub fn from_backing_pressure(
2055        observed: ExecutorAdmissionEpochs,
2056        wait_condition: crate::vnext::CapacityWaitCondition,
2057        pressure: crate::vnext::DynamicBackingPressure,
2058        stage: ExecutorExecutionCapacityStage,
2059    ) -> Result<Self> {
2060        Self::with_evidence(
2061            observed,
2062            wait_condition,
2063            stage,
2064            ExecutorExecutionCapacityEvidence::direct_backing_pressure(pressure),
2065        )
2066    }
2067
2068    pub fn from_admission(
2069        deferred: &crate::vnext::AdmissionDeferred,
2070        stage: ExecutorExecutionCapacityStage,
2071    ) -> Result<Self> {
2072        if deferred.action() != crate::vnext::DeferredAction::WaitForRelease {
2073            return Err(ferrum_types::FerrumError::internal(
2074                "execution capacity deferral must be reduced to WaitForRelease before export",
2075            ));
2076        }
2077        let evidence = ExecutorExecutionCapacityEvidence::logical(deferred.blockers().to_vec())?;
2078        Self::with_evidence(
2079            ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2080            deferred.wait_condition().clone(),
2081            stage,
2082            evidence,
2083        )
2084    }
2085
2086    /// Export an unresolved logical backing-growth decision after the
2087    /// executor's bounded in-call maintenance attempts are exhausted.
2088    ///
2089    /// This remains a pre-submit scheduling deferral. The bound controls only
2090    /// how much allocator maintenance one executor call may perform; it must
2091    /// not turn temporary capacity pressure into a terminal request failure.
2092    pub fn from_pending_maintenance(
2093        deferred: &crate::vnext::AdmissionDeferred,
2094        stage: ExecutorExecutionCapacityStage,
2095    ) -> Result<Self> {
2096        if deferred.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2097            return Err(ferrum_types::FerrumError::internal(
2098                "pending execution maintenance must await backing growth",
2099            ));
2100        }
2101        let evidence = ExecutorExecutionCapacityEvidence::logical(deferred.blockers().to_vec())?;
2102        Self::with_evidence(
2103            ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2104            deferred.wait_condition().clone(),
2105            stage,
2106            evidence,
2107        )
2108    }
2109
2110    /// Export unresolved physical backing pressure without discarding its
2111    /// exact pool/domain evidence.
2112    pub fn from_backing(
2113        deferred: &crate::vnext::DynamicBackingDeferred,
2114        stage: ExecutorExecutionCapacityStage,
2115    ) -> Result<Self> {
2116        let evidence =
2117            ExecutorExecutionCapacityEvidence::backing_deferred(deferred.blockers().to_vec())?;
2118        Self::with_evidence(
2119            ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2120            deferred.wait_condition().clone(),
2121            stage,
2122            evidence,
2123        )
2124    }
2125
2126    /// Attach a scheduler retry only when this call produced at least one real
2127    /// physical mutation relevant to the final blocker. An empty or unrelated
2128    /// receipt set is an ordinary typed wait, not an internal error.
2129    pub fn with_relevant_maintenance_retry(
2130        mut self,
2131        attempts: u32,
2132        receipts: &[crate::vnext::DynamicPoolGrowthBatchReceipt],
2133        pools: &[crate::vnext::DynamicPoolStatus],
2134        affected_request_ids: Vec<RequestId>,
2135    ) -> Result<Self> {
2136        if receipts.is_empty() {
2137            return Ok(self);
2138        }
2139        let progress = ExecutorExecutionMaintenanceProgress::from_growth_receipts(
2140            attempts,
2141            self.observed,
2142            receipts,
2143            pools,
2144        )?;
2145        if progress.coordinator_id() != self.observed.coordinator_id
2146            || progress.latest_capacity_epoch() > self.observed.capacity_epoch
2147            || progress.mutations().is_empty()
2148        {
2149            return Err(FerrumError::internal(
2150                "execution maintenance progress does not match the exported deferral",
2151            ));
2152        }
2153        let relevant_mutation = progress
2154            .mutations()
2155            .iter()
2156            .any(|mutation| self.evidence.has_relevant_mutation(mutation));
2157        if !relevant_mutation {
2158            return Ok(self);
2159        }
2160        let retry = ExecutorExecutionMaintenanceRetry::new(affected_request_ids, progress)?;
2161        if self.stage == ExecutorExecutionCapacityStage::SequenceExtension
2162            && retry.affected_request_ids().len() != 1
2163        {
2164            return Err(FerrumError::internal(
2165                "sequence-extension maintenance retry must affect exactly one request",
2166            ));
2167        }
2168        self.maintenance_retry = Some(retry);
2169        Ok(self)
2170    }
2171
2172    pub fn from_admission_maintenance(
2173        source: &crate::vnext::AdmissionDeferred,
2174        observed: ExecutorAdmissionEpochs,
2175        wait_condition: crate::vnext::CapacityWaitCondition,
2176        pressure: crate::vnext::DynamicBackingPressure,
2177        maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
2178        stage: ExecutorExecutionCapacityStage,
2179    ) -> Result<Self> {
2180        if source.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2181            return Err(ferrum_types::FerrumError::internal(
2182                "execution maintenance source must await backing growth",
2183            ));
2184        }
2185        let evidence = ExecutorExecutionCapacityEvidence::logical_with_pressure(
2186            source.blockers().to_vec(),
2187            Some(pressure),
2188        )?
2189        .with_maintenance_boundary(maintenance_boundary)?;
2190        if evidence.maintenance_boundary().is_some_and(|boundary| {
2191            boundary.coordinator_id().get() != observed.coordinator_id.get()
2192        }) {
2193            return Err(FerrumError::internal(
2194                "execution maintenance boundary belongs to another coordinator",
2195            ));
2196        }
2197        Self::with_evidence(observed, wait_condition, stage, evidence)
2198    }
2199
2200    pub fn from_backing_maintenance(
2201        source: &crate::vnext::DynamicBackingDeferred,
2202        observed: ExecutorAdmissionEpochs,
2203        wait_condition: crate::vnext::CapacityWaitCondition,
2204        pressure: crate::vnext::DynamicBackingPressure,
2205        maintenance_boundary: Option<crate::vnext::DynamicPoolMaintenanceBoundaryReceipt>,
2206        stage: ExecutorExecutionCapacityStage,
2207    ) -> Result<Self> {
2208        let evidence = ExecutorExecutionCapacityEvidence::backing_deferred_with_pressure(
2209            source.blockers().to_vec(),
2210            Some(pressure),
2211        )?
2212        .with_maintenance_boundary(maintenance_boundary)?;
2213        if evidence.maintenance_boundary().is_some_and(|boundary| {
2214            boundary.coordinator_id().get() != observed.coordinator_id.get()
2215        }) {
2216            return Err(FerrumError::internal(
2217                "execution maintenance boundary belongs to another coordinator",
2218            ));
2219        }
2220        Self::with_evidence(observed, wait_condition, stage, evidence)
2221    }
2222
2223    pub const fn observed(&self) -> ExecutorAdmissionEpochs {
2224        self.observed
2225    }
2226
2227    pub fn wait_condition(&self) -> &crate::vnext::CapacityWaitCondition {
2228        &self.wait_condition
2229    }
2230
2231    pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2232        self.stage
2233    }
2234
2235    pub const fn evidence(&self) -> &ExecutorExecutionCapacityEvidence {
2236        &self.evidence
2237    }
2238
2239    pub fn shortfalls(&self) -> &[crate::vnext::CapacityShortfall] {
2240        self.evidence.shortfalls()
2241    }
2242
2243    pub fn backing_blockers(&self) -> &[crate::vnext::DynamicBackingBlocker] {
2244        self.evidence.backing_blockers()
2245    }
2246
2247    pub const fn backing_pressure(&self) -> Option<&crate::vnext::DynamicBackingPressure> {
2248        self.evidence.backing_pressure()
2249    }
2250
2251    pub const fn maintenance_boundary(
2252        &self,
2253    ) -> Option<&crate::vnext::DynamicPoolMaintenanceBoundaryReceipt> {
2254        self.evidence.maintenance_boundary()
2255    }
2256
2257    pub fn maintenance_retry(&self) -> Option<&ExecutorExecutionMaintenanceRetry> {
2258        self.maintenance_retry.as_ref()
2259    }
2260
2261    /// Validate the bound retry scope against the authoritative logical
2262    /// frontiers in the current engine call.
2263    pub fn validated_maintenance_retry_scope(
2264        &self,
2265        current_request_ids: &[RequestId],
2266    ) -> Result<Option<&ExecutorExecutionMaintenanceRetry>> {
2267        let Some(retry) = self.maintenance_retry.as_ref() else {
2268            return Ok(None);
2269        };
2270        let current = current_request_ids.iter().collect::<HashSet<_>>();
2271        if current_request_ids.is_empty() || current.len() != current_request_ids.len() {
2272            return Err(FerrumError::internal(
2273                "execution maintenance retry received an invalid current request cohort",
2274            ));
2275        }
2276        let affected = retry.affected_request_ids().iter().collect::<HashSet<_>>();
2277        if !affected.is_subset(&current) {
2278            return Err(FerrumError::internal(
2279                "execution maintenance retry affects a request outside the current cohort",
2280            ));
2281        }
2282        match self.stage {
2283            ExecutorExecutionCapacityStage::SequenceExtension => {
2284                if affected.len() != 1 {
2285                    return Err(FerrumError::internal(
2286                        "sequence-extension maintenance retry must affect one current request",
2287                    ));
2288                }
2289            }
2290            ExecutorExecutionCapacityStage::StepAdmission
2291            | ExecutorExecutionCapacityStage::SubmissionWave => {
2292                if affected != current {
2293                    return Err(FerrumError::internal(
2294                        "cohort maintenance retry must cover the complete current cohort",
2295                    ));
2296                }
2297            }
2298        }
2299        Ok(Some(retry))
2300    }
2301
2302    /// Return a strictly smaller, capacity-informed prefill width.
2303    ///
2304    /// This is a cold pressure-path hint, not allocator authority. The caller
2305    /// must probe the returned prefix through normal typed admission before any
2306    /// provider encode or device submission. A bounded reduction prevents a
2307    /// large frontier from producing an unbounded sequence of near-identical
2308    /// probes when the shortfall is small.
2309    pub fn narrower_prefill_tokens(&self, attempted_tokens: usize) -> Option<usize> {
2310        if attempted_tokens <= 1 {
2311            return None;
2312        }
2313        let maximum_next = attempted_tokens
2314            .saturating_sub(attempted_tokens.div_ceil(4))
2315            .max(1);
2316        let proportional = self
2317            .shortfalls()
2318            .iter()
2319            .filter_map(|shortfall| {
2320                let requested = shortfall.requested().get();
2321                let available = shortfall.available().get();
2322                (requested > available).then(|| {
2323                    let scaled = (attempted_tokens as u128).saturating_mul(available as u128)
2324                        / requested as u128;
2325                    usize::try_from(scaled)
2326                        .unwrap_or(usize::MAX)
2327                        .clamp(1, attempted_tokens - 1)
2328                })
2329            })
2330            .min();
2331        Some(
2332            proportional
2333                .unwrap_or_else(|| attempted_tokens.div_ceil(2))
2334                .min(maximum_next)
2335                .max(1),
2336        )
2337    }
2338}
2339
2340/// Scheduler-visible proof that execution is temporarily blocked by a
2341/// Request-lifetime state hazard rather than by physical capacity.
2342///
2343/// The embedded hazard evidence retains the exact plan-local coordinator and
2344/// supports subscribe-before-recheck waiter registration. The request cohort
2345/// is the product identity projection used by the scheduler; it must stay
2346/// separate from the allocator's internal request-authority ids.
2347#[derive(Debug, Clone, Serialize)]
2348pub struct ExecutorRequestStateDeferral {
2349    stage: ExecutorExecutionCapacityStage,
2350    request_ids: Vec<RequestId>,
2351    hazard: crate::vnext::RequestStateHazardDeferral,
2352}
2353
2354impl ExecutorRequestStateDeferral {
2355    pub fn new(
2356        stage: ExecutorExecutionCapacityStage,
2357        request_ids: Vec<RequestId>,
2358        hazard: crate::vnext::RequestStateHazardDeferral,
2359    ) -> Result<Self> {
2360        let unique = request_ids.iter().collect::<HashSet<_>>();
2361        if request_ids.is_empty() || unique.len() != request_ids.len() {
2362            return Err(FerrumError::internal(
2363                "request-state execution deferral requires a non-empty unique product cohort",
2364            ));
2365        }
2366        if hazard.blockers().is_empty() {
2367            return Err(FerrumError::internal(
2368                "request-state execution deferral requires exact blockers",
2369            ));
2370        }
2371        Ok(Self {
2372            stage,
2373            request_ids,
2374            hazard,
2375        })
2376    }
2377
2378    pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2379        self.stage
2380    }
2381
2382    pub fn request_ids(&self) -> &[RequestId] {
2383        &self.request_ids
2384    }
2385
2386    pub const fn hazard(&self) -> &crate::vnext::RequestStateHazardDeferral {
2387        &self.hazard
2388    }
2389
2390    pub fn register_waiter(&self) -> Result<crate::vnext::RequestStateHazardWaitRegistration> {
2391        self.hazard
2392            .register_waiter()
2393            .map_err(|error| FerrumError::backend(error.to_string()))
2394    }
2395}
2396
2397/// A pre-submit execution frontier can be blocked by independently evolving
2398/// sources. Capacity waits participate in scheduler pressure/yield policy;
2399/// Request-state waits never do and resume only from their exact hazard
2400/// coordinator.
2401#[derive(Debug, Clone, Serialize)]
2402#[serde(tag = "reason", content = "evidence", rename_all = "snake_case")]
2403pub enum ExecutorExecutionDeferral {
2404    Capacity(ExecutorExecutionCapacityDeferral),
2405    RequestState(ExecutorRequestStateDeferral),
2406}
2407
2408impl ExecutorExecutionDeferral {
2409    pub const fn stage(&self) -> ExecutorExecutionCapacityStage {
2410        match self {
2411            Self::Capacity(deferral) => deferral.stage(),
2412            Self::RequestState(deferral) => deferral.stage(),
2413        }
2414    }
2415
2416    pub const fn as_capacity(&self) -> Option<&ExecutorExecutionCapacityDeferral> {
2417        match self {
2418            Self::Capacity(deferral) => Some(deferral),
2419            Self::RequestState(_) => None,
2420        }
2421    }
2422
2423    pub const fn as_request_state(&self) -> Option<&ExecutorRequestStateDeferral> {
2424        match self {
2425            Self::Capacity(_) => None,
2426            Self::RequestState(deferral) => Some(deferral),
2427        }
2428    }
2429}
2430
2431impl From<ExecutorExecutionCapacityDeferral> for ExecutorExecutionDeferral {
2432    fn from(deferral: ExecutorExecutionCapacityDeferral) -> Self {
2433        Self::Capacity(deferral)
2434    }
2435}
2436
2437impl From<ExecutorRequestStateDeferral> for ExecutorExecutionDeferral {
2438    fn from(deferral: ExecutorRequestStateDeferral) -> Self {
2439        Self::RequestState(deferral)
2440    }
2441}
2442
2443#[cfg(test)]
2444mod execution_capacity_deferral_tests {
2445    use super::{
2446        ExecutorAdmissionEpochs, ExecutorExecutionCapacityDeferral,
2447        ExecutorExecutionCapacityEvidenceOwner, ExecutorExecutionCapacityStage,
2448        ExecutorExecutionMaintenanceProgress, ExecutorExecutionMaintenanceRetry,
2449    };
2450    use crate::vnext::{
2451        CapacityAvailabilityEpoch, CapacityAvailabilitySource, CapacityWaitCondition,
2452        DeviceCapacityPressure, DeviceCapacityPressureScope, DynamicBackingPressure,
2453    };
2454    use ferrum_types::RequestId;
2455    use std::num::NonZeroU64;
2456
2457    fn test_progress() -> ExecutorExecutionMaintenanceProgress {
2458        ExecutorExecutionMaintenanceProgress {
2459            attempts: 1,
2460            coordinator_id: NonZeroU64::new(19).unwrap(),
2461            mutations: Vec::new(),
2462            latest_capacity_epoch: 5,
2463        }
2464    }
2465
2466    fn test_deferral(stage: ExecutorExecutionCapacityStage) -> ExecutorExecutionCapacityDeferral {
2467        let observed =
2468            CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 7)
2469                .unwrap();
2470        let condition = CapacityWaitCondition::from_observation(19, vec![observed]).unwrap();
2471        ExecutorExecutionCapacityDeferral::from_backing_pressure(
2472            ExecutorAdmissionEpochs::new(NonZeroU64::new(19).unwrap(), 3, 5),
2473            condition,
2474            test_pressure(),
2475            stage,
2476        )
2477        .unwrap()
2478    }
2479
2480    fn test_pressure() -> DynamicBackingPressure {
2481        DeviceCapacityPressure::new(
2482            DeviceCapacityPressureScope::PlanBudget,
2483            "device.execution-capacity-test".to_owned(),
2484            1,
2485            1,
2486            1,
2487            1,
2488            1,
2489        )
2490        .unwrap()
2491        .into()
2492    }
2493
2494    #[test]
2495    fn prefill_narrowing_is_strict_bounded_and_stops_at_one_token() {
2496        let observed =
2497            CapacityAvailabilityEpoch::new(CapacityAvailabilitySource::ActiveSequenceSlots, 7)
2498                .unwrap();
2499        let condition = CapacityWaitCondition::from_observation(19, vec![observed]).unwrap();
2500        let deferred = ExecutorExecutionCapacityDeferral::from_backing_pressure(
2501            ExecutorAdmissionEpochs::new(NonZeroU64::new(19).unwrap(), 3, 5),
2502            condition,
2503            test_pressure(),
2504            ExecutorExecutionCapacityStage::StepAdmission,
2505        )
2506        .unwrap();
2507
2508        assert_eq!(deferred.narrower_prefill_tokens(342), Some(171));
2509        assert_eq!(deferred.narrower_prefill_tokens(2), Some(1));
2510        assert_eq!(deferred.narrower_prefill_tokens(1), None);
2511    }
2512
2513    #[test]
2514    fn backing_pressure_serializes_one_typed_evidence_owner() {
2515        let deferred = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension);
2516        let serialized = serde_json::to_value(&deferred).unwrap();
2517
2518        assert_eq!(
2519            deferred.evidence().owner(),
2520            ExecutorExecutionCapacityEvidenceOwner::Backing
2521        );
2522        assert!(deferred.shortfalls().is_empty());
2523        assert!(deferred.backing_blockers().is_empty());
2524        assert!(deferred.backing_pressure().is_some());
2525        assert_eq!(serialized["evidence"]["owner"], "backing");
2526        assert_eq!(serialized["evidence"]["kind"], "backing_pressure");
2527        assert!(serialized["evidence"]["pressure"].is_object());
2528    }
2529
2530    #[test]
2531    fn empty_maintenance_receipts_remain_an_ordinary_typed_deferral() {
2532        let request_id = RequestId::new();
2533        let deferred = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension)
2534            .with_relevant_maintenance_retry(2, &[], &[], vec![request_id])
2535            .unwrap();
2536
2537        assert!(deferred.maintenance_retry().is_none());
2538    }
2539
2540    #[test]
2541    fn maintenance_retry_rejects_empty_duplicate_or_unproven_scope() {
2542        let request_id = RequestId::new();
2543        assert!(ExecutorExecutionMaintenanceRetry::new(Vec::new(), test_progress()).is_err());
2544        assert!(ExecutorExecutionMaintenanceRetry::new(
2545            vec![request_id.clone(), request_id.clone()],
2546            test_progress(),
2547        )
2548        .is_err());
2549        assert!(ExecutorExecutionMaintenanceRetry::new(vec![request_id], test_progress()).is_err());
2550    }
2551
2552    #[test]
2553    fn maintenance_retry_scope_is_fail_closed_for_sequence_and_cohort_stages() {
2554        let first = RequestId::new();
2555        let second = RequestId::new();
2556        let retry = |affected_request_ids| ExecutorExecutionMaintenanceRetry {
2557            affected_request_ids,
2558            progress: test_progress(),
2559        };
2560
2561        let mut sequence = test_deferral(ExecutorExecutionCapacityStage::SequenceExtension);
2562        sequence.maintenance_retry = Some(retry(vec![second.clone()]));
2563        assert_eq!(
2564            sequence
2565                .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2566                .unwrap()
2567                .unwrap()
2568                .affected_request_ids(),
2569            [second.clone()]
2570        );
2571        sequence.maintenance_retry = Some(retry(vec![first.clone(), second.clone()]));
2572        assert!(sequence
2573            .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2574            .is_err());
2575
2576        let mut cohort = test_deferral(ExecutorExecutionCapacityStage::SubmissionWave);
2577        cohort.maintenance_retry = Some(retry(vec![second.clone()]));
2578        assert!(cohort
2579            .validated_maintenance_retry_scope(&[first.clone(), second.clone()])
2580            .is_err());
2581        cohort.maintenance_retry = Some(retry(vec![first.clone(), second.clone()]));
2582        assert!(cohort
2583            .validated_maintenance_retry_scope(&[first, second])
2584            .unwrap()
2585            .is_some());
2586    }
2587}
2588
2589/// Capacity-aware batch decode result.
2590///
2591/// `Deferred` is only legal before provider encode or device submission. All
2592/// possibly-submitted failures remain ordinary errors and retain their typed
2593/// fence/recovery authority inside the executor.
2594pub enum ExecutorBatchDecodeOutcome {
2595    Completed(Vec<DecodeOutput>),
2596    Deferred(ExecutorExecutionDeferral),
2597}
2598
2599/// Capacity-aware, tensor-free batch decode result for a plan runtime.
2600pub enum PlanRuntimeBatchDecodeOutcome {
2601    Completed(Vec<PlanRuntimeDecodeOutput>),
2602    Deferred(ExecutorExecutionDeferral),
2603}
2604
2605/// Capacity-aware result for one tensor-free plan-runtime prefill frontier.
2606pub struct PlanRuntimePrefillCompletion {
2607    output: PlanRuntimePrefillOutput,
2608    planned_chunk: PrefillChunk,
2609    completed_chunk: PrefillChunk,
2610    capacity_probe_count: u32,
2611}
2612
2613impl PlanRuntimePrefillCompletion {
2614    pub fn new(
2615        output: PlanRuntimePrefillOutput,
2616        planned_chunk: PrefillChunk,
2617        completed_chunk: PrefillChunk,
2618        capacity_probe_count: u32,
2619    ) -> Result<Self> {
2620        validate_prefill_completion_shape(planned_chunk, completed_chunk, capacity_probe_count)?;
2621        Ok(Self {
2622            output,
2623            planned_chunk,
2624            completed_chunk,
2625            capacity_probe_count,
2626        })
2627    }
2628
2629    pub fn exact(output: PlanRuntimePrefillOutput, chunk: PrefillChunk) -> Self {
2630        Self {
2631            output,
2632            planned_chunk: chunk,
2633            completed_chunk: chunk,
2634            capacity_probe_count: 0,
2635        }
2636    }
2637
2638    pub const fn planned_chunk(&self) -> PrefillChunk {
2639        self.planned_chunk
2640    }
2641
2642    pub const fn completed_chunk(&self) -> PrefillChunk {
2643        self.completed_chunk
2644    }
2645
2646    pub const fn capacity_probe_count(&self) -> u32 {
2647        self.capacity_probe_count
2648    }
2649
2650    pub fn output(&self) -> &PlanRuntimePrefillOutput {
2651        &self.output
2652    }
2653
2654    pub fn validate_for(
2655        &self,
2656        expected_request_id: &RequestId,
2657        expected_planned_chunk: PrefillChunk,
2658        vocabulary_size: usize,
2659    ) -> Result<()> {
2660        if self.planned_chunk != expected_planned_chunk {
2661            return Err(FerrumError::backend(format!(
2662                "plan runtime completed prefill frontier {:?}, expected {:?}",
2663                self.planned_chunk.range(),
2664                expected_planned_chunk.range()
2665            )));
2666        }
2667        validate_prefill_completion_shape(
2668            self.planned_chunk,
2669            self.completed_chunk,
2670            self.capacity_probe_count,
2671        )?;
2672        self.output.validate_for_completion(
2673            expected_request_id,
2674            self.completed_chunk,
2675            vocabulary_size,
2676        )
2677    }
2678
2679    pub fn into_parts(self) -> (PlanRuntimePrefillOutput, PrefillChunk, PrefillChunk, u32) {
2680        (
2681            self.output,
2682            self.planned_chunk,
2683            self.completed_chunk,
2684            self.capacity_probe_count,
2685        )
2686    }
2687}
2688
2689pub enum PlanRuntimePrefillOutcome {
2690    Completed(PlanRuntimePrefillCompletion),
2691    Deferred(ExecutorExecutionDeferral),
2692}
2693
2694pub enum PlanRuntimeBatchPrefillOutcome {
2695    Completed(Vec<PlanRuntimePrefillCompletion>),
2696    NotSubmitted(ExecutorExecutionDeferral),
2697    Unsupported,
2698}
2699
2700/// Capacity-aware result for one planned prefill frontier.
2701///
2702/// `Deferred` is only legal before provider encode or device submission. The
2703/// executor may complete a strict prefix after typed capacity probes; the
2704/// scheduler commits only that prefix and learns the narrower execution ceiling.
2705pub struct ExecutorPrefillCompletion {
2706    output: PrefillOutput,
2707    planned_chunk: PrefillChunk,
2708    completed_chunk: PrefillChunk,
2709    capacity_probe_count: u32,
2710}
2711
2712impl ExecutorPrefillCompletion {
2713    pub fn new(
2714        output: PrefillOutput,
2715        planned_chunk: PrefillChunk,
2716        completed_chunk: PrefillChunk,
2717        capacity_probe_count: u32,
2718    ) -> Result<Self> {
2719        validate_prefill_completion_shape(planned_chunk, completed_chunk, capacity_probe_count)?;
2720        Ok(Self {
2721            output,
2722            planned_chunk,
2723            completed_chunk,
2724            capacity_probe_count,
2725        })
2726    }
2727
2728    pub fn exact(output: PrefillOutput, chunk: PrefillChunk) -> Self {
2729        Self {
2730            output,
2731            planned_chunk: chunk,
2732            completed_chunk: chunk,
2733            capacity_probe_count: 0,
2734        }
2735    }
2736
2737    pub const fn planned_chunk(&self) -> PrefillChunk {
2738        self.planned_chunk
2739    }
2740
2741    pub const fn completed_chunk(&self) -> PrefillChunk {
2742        self.completed_chunk
2743    }
2744
2745    pub const fn capacity_probe_count(&self) -> u32 {
2746        self.capacity_probe_count
2747    }
2748
2749    pub fn into_parts(self) -> (PrefillOutput, PrefillChunk, PrefillChunk, u32) {
2750        (
2751            self.output,
2752            self.planned_chunk,
2753            self.completed_chunk,
2754            self.capacity_probe_count,
2755        )
2756    }
2757}
2758
2759fn validate_prefill_completion_shape(
2760    planned_chunk: PrefillChunk,
2761    completed_chunk: PrefillChunk,
2762    capacity_probe_count: u32,
2763) -> Result<()> {
2764    if completed_chunk.tokens_processed() != planned_chunk.tokens_processed()
2765        || completed_chunk.total_prompt_tokens() != planned_chunk.total_prompt_tokens()
2766        || completed_chunk.tokens_to_process() > planned_chunk.tokens_to_process()
2767    {
2768        return Err(ferrum_types::FerrumError::internal(
2769            "completed prefill chunk is not a non-empty prefix of its planned chunk",
2770        ));
2771    }
2772    if completed_chunk != planned_chunk && capacity_probe_count == 0 {
2773        return Err(ferrum_types::FerrumError::internal(
2774            "partial prefill completion requires a failed capacity probe",
2775        ));
2776    }
2777    Ok(())
2778}
2779
2780pub enum ExecutorPrefillOutcome {
2781    Completed(ExecutorPrefillCompletion),
2782    Deferred(ExecutorExecutionDeferral),
2783}
2784
2785/// Transactional result of attempting one physical prefill batch.
2786///
2787/// `NotSubmitted` proves that no participant in the batch reached provider
2788/// encode or device submission. The caller may therefore retry a narrower
2789/// partition or the existing per-request capacity path without duplicating
2790/// model work. `Unsupported` keeps the optimization optional for legacy
2791/// executors while plan-runtime implementations provide the real batch edge.
2792pub enum ExecutorBatchPrefillOutcome {
2793    Completed(Vec<ExecutorPrefillCompletion>),
2794    NotSubmitted(ExecutorExecutionDeferral),
2795    Unsupported,
2796}
2797
2798/// Stage that must advance before a plan-runtime prefill can be admitted.
2799///
2800/// This is scheduler evidence, not allocator authority. The executor retains
2801/// the sealed logical or physical deferral that authorizes maintenance.
2802#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
2803#[serde(rename_all = "snake_case")]
2804pub enum ExecutorPrefillMaintenanceStage {
2805    LogicalCapacity,
2806    PhysicalBacking,
2807}
2808
2809/// Scheduler-visible reason that a prefill needs plan-runtime backing
2810/// maintenance. These values are projections only and cannot allocate memory.
2811#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2812#[serde(tag = "source", rename_all = "snake_case")]
2813pub enum ExecutorPrefillMaintenanceBlocker {
2814    Capacity {
2815        domain_id: Option<u32>,
2816        kind: crate::vnext::CapacityShortfallKind,
2817        requested: u64,
2818        available: u64,
2819        current_total: u64,
2820        maximum_total: u64,
2821    },
2822    Backing {
2823        pool_id: String,
2824        domain_id: u32,
2825        lifetime: crate::vnext::DynamicBackingClaimScope,
2826        reason: crate::vnext::DynamicBackingDeferralReason,
2827        requested_bytes: u64,
2828        free_bytes: u64,
2829        largest_contiguous_bytes: u64,
2830    },
2831}
2832
2833/// Non-authoritative projection of plan-runtime maintenance work.
2834///
2835/// The request id is the only handle returned to the engine. Implementations
2836/// must retain the sealed deferral internally and validate it again when
2837/// [`ModelExecutor::maintain_prefill_backing`] is called.
2838#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2839pub struct ExecutorPrefillMaintenanceDeferral {
2840    request_id: RequestId,
2841    observed: ExecutorAdmissionEpochs,
2842    wait_condition: crate::vnext::CapacityWaitCondition,
2843    stage: ExecutorPrefillMaintenanceStage,
2844    blockers: Vec<ExecutorPrefillMaintenanceBlocker>,
2845}
2846
2847impl ExecutorPrefillMaintenanceDeferral {
2848    pub fn new(
2849        request_id: RequestId,
2850        observed: ExecutorAdmissionEpochs,
2851        wait_condition: crate::vnext::CapacityWaitCondition,
2852        stage: ExecutorPrefillMaintenanceStage,
2853        blockers: Vec<ExecutorPrefillMaintenanceBlocker>,
2854    ) -> Result<Self> {
2855        if blockers.is_empty() {
2856            return Err(ferrum_types::FerrumError::request_validation(
2857                "executor prefill maintenance deferral requires at least one blocker",
2858            ));
2859        }
2860        if wait_condition.coordinator_id().get() != observed.coordinator_id.get() {
2861            return Err(ferrum_types::FerrumError::request_validation(
2862                "executor prefill maintenance wait condition belongs to a different coordinator",
2863            ));
2864        }
2865        Ok(Self {
2866            request_id,
2867            observed,
2868            wait_condition,
2869            stage,
2870            blockers,
2871        })
2872    }
2873
2874    pub fn from_admission(
2875        request_id: &RequestId,
2876        deferred: &crate::vnext::AdmissionDeferred,
2877    ) -> Result<Self> {
2878        if deferred.action() != crate::vnext::DeferredAction::AwaitBackingGrowth {
2879            return Err(ferrum_types::FerrumError::internal(
2880                "logical prefill maintenance projection requires AwaitBackingGrowth",
2881            ));
2882        }
2883        let blockers = deferred
2884            .blockers()
2885            .iter()
2886            .map(|blocker| ExecutorPrefillMaintenanceBlocker::Capacity {
2887                domain_id: blocker.domain().map(|domain| domain.get()),
2888                kind: blocker.kind(),
2889                requested: blocker.requested().get(),
2890                available: blocker.available().get(),
2891                current_total: blocker.current_total().get(),
2892                maximum_total: blocker.maximum_total().get(),
2893            })
2894            .collect();
2895        Self::new(
2896            request_id.clone(),
2897            ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2898            deferred.wait_condition().clone(),
2899            ExecutorPrefillMaintenanceStage::LogicalCapacity,
2900            blockers,
2901        )
2902    }
2903
2904    pub fn from_backing(
2905        request_id: &RequestId,
2906        deferred: &crate::vnext::DynamicBackingDeferred,
2907    ) -> Result<Self> {
2908        let blockers = deferred
2909            .blockers()
2910            .iter()
2911            .map(|blocker| ExecutorPrefillMaintenanceBlocker::Backing {
2912                pool_id: blocker.pool_id().as_str().to_string(),
2913                domain_id: blocker.domain_id().get(),
2914                lifetime: deferred.scope(),
2915                reason: blocker.reason(),
2916                requested_bytes: blocker.requested_bytes(),
2917                free_bytes: blocker.free_bytes(),
2918                largest_contiguous_bytes: blocker.largest_contiguous_bytes(),
2919            })
2920            .collect();
2921        Self::new(
2922            request_id.clone(),
2923            ExecutorAdmissionEpochs::from_capacity(deferred.epochs()),
2924            deferred.wait_condition().clone(),
2925            ExecutorPrefillMaintenanceStage::PhysicalBacking,
2926            blockers,
2927        )
2928    }
2929
2930    pub fn request_id(&self) -> &RequestId {
2931        &self.request_id
2932    }
2933
2934    pub const fn observed(&self) -> ExecutorAdmissionEpochs {
2935        self.observed
2936    }
2937
2938    pub fn wait_condition(&self) -> &crate::vnext::CapacityWaitCondition {
2939        &self.wait_condition
2940    }
2941
2942    pub const fn stage(&self) -> ExecutorPrefillMaintenanceStage {
2943        self.stage
2944    }
2945
2946    pub fn blockers(&self) -> &[ExecutorPrefillMaintenanceBlocker] {
2947        &self.blockers
2948    }
2949}
2950
2951/// Result of one bounded plan-runtime backing maintenance attempt.
2952#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
2953#[serde(tag = "outcome", rename_all = "snake_case")]
2954pub enum ExecutorPrefillMaintenanceOutcome {
2955    /// Cancellation won the race before the maintenance task consumed the
2956    /// retained deferral.
2957    NoLongerPending,
2958    /// The physical allocator changed while maintenance was installing its
2959    /// wait predicate. The scheduler must clear the old backing deferral and
2960    /// perform one authoritative admission probe, even if publication of the
2961    /// corresponding capacity epoch is still in flight.
2962    RetryAdmission { current: ExecutorAdmissionEpochs },
2963    /// The requested backing is valid but cannot be installed while current
2964    /// device claims remain live. The scheduler must wait for release evidence
2965    /// rather than completing the request as an error.
2966    WaitForRelease {
2967        current: ExecutorAdmissionEpochs,
2968        wait_condition: crate::vnext::CapacityWaitCondition,
2969        pressure: crate::vnext::DynamicBackingPressure,
2970    },
2971    /// The executor installed real backing and published the resulting
2972    /// capacity epoch.
2973    Maintained {
2974        current: ExecutorAdmissionEpochs,
2975        pools_grown: usize,
2976        allocated_bytes: u64,
2977        pools_reclaimed: usize,
2978        chunks_reclaimed: usize,
2979        reclaimed_bytes: u64,
2980        /// Exact allocator-issued rebalance receipt. The aggregate counters
2981        /// above remain for stable metrics and must reconcile with this value.
2982        rebalance: Option<crate::vnext::DynamicPoolRebalanceReceipt>,
2983    },
2984}
2985
2986/// Typed result of probing plan-runtime prefill capacity.
2987///
2988/// `Deferred` and `MaintenanceDeferred` preserve the capacity domains and
2989/// epochs required by the plan-local dynamic admission queue. They must never
2990/// be flattened into a generic resource error at the scheduler boundary.
2991#[derive(Debug, Clone)]
2992pub enum ExecutorPrefillAdmissionDecision {
2993    Admitted(ExecutorPrefillAdmissionReceipt),
2994    Deferred(crate::vnext::AdmissionDeferred),
2995    MaintenanceDeferred(ExecutorPrefillMaintenanceDeferral),
2996    PermanentRejected(crate::vnext::AdmissionRejected),
2997}
2998
2999/// Core model executor trait focusing on tensor operations
3000#[async_trait]
3001pub trait ModelExecutor: Send + Sync {
3002    /// Plan an optional prompt-tail checkpoint before a prefill chunk is
3003    /// dispatched. The boundary must lie after this chunk's start and no later
3004    /// than its end, leaving a legal suffix for logits. None preserves the chunk. Planning
3005    /// grants no capacity or capture authority and does not report a failed
3006    /// capacity probe; the caller dispatches the resulting chunk explicitly.
3007    fn plan_prompt_tail_capture_boundary(&self, _chunk: PrefillChunk) -> Option<PrefixCapturePlan> {
3008        None
3009    }
3010
3011    /// Pure boundary planning for optional sharing. None leaves normal scheduling
3012    /// unchanged; a returned boundary never grants request or resource authority.
3013    fn plan_prefix_capture_boundary(
3014        &self,
3015        _input: PrefixCaptureBoundary<'_>,
3016    ) -> Option<PrefixCapturePlan> {
3017        None
3018    }
3019
3020    /// Arm interest only against an already-admitted, exact source incarnation.
3021    /// No device allocation, capacity reservation, provider encoding, or device
3022    /// submission is allowed here.
3023    fn retain_prefix_capture_interest(
3024        &self,
3025        _input: PrefixCaptureRequest<'_>,
3026    ) -> Result<Option<Arc<dyn PrefixCaptureLease>>> {
3027        Ok(None)
3028    }
3029
3030    /// Get model information and metadata
3031    fn info(&self) -> &ModelInfo;
3032
3033    /// Selects the single authority for request-lifetime model resources.
3034    /// Existing executors remain on the transitional legacy-engine path by
3035    /// default. A runtime that returns `PlanRuntime` must return the shared
3036    /// runtime's opaque cache handle from prefill/decode and delegate release
3037    /// of that authority from `release_cache`.
3038    fn execution_resource_authority(&self) -> ExecutionResourceAuthority {
3039        ExecutionResourceAuthority::LegacyEngine
3040    }
3041
3042    /// Immutable admission limits compiled into this executor's runtime plan.
3043    /// A PlanRuntime executor must return `Some`; legacy executors may defer to
3044    /// the engine-owned scheduler and recurrent-state limits.
3045    fn admission_limits(&self) -> Result<Option<ExecutorAdmissionLimits>> {
3046        Ok(None)
3047    }
3048
3049    /// Returns the immutable product plan that owns planning, provider
3050    /// selection, and resource authority for a plan-runtime executor.
3051    /// Legacy executors return `None`; `PlanRuntime` executors must expose the
3052    /// exact plan used for provisioning and dispatch.
3053    fn resolved_model_plan(&self) -> Option<&crate::vnext::ResolvedModelPlan> {
3054        None
3055    }
3056
3057    /// Returns the authoritative memory breakdown for a shared plan runtime.
3058    /// `LegacyEngine` executors return `None`; `PlanRuntime` executors must
3059    /// return `Some` while they are ready.
3060    fn plan_runtime_resource_snapshot(&self) -> Result<Option<PlanRuntimeResourceSnapshot>> {
3061        Ok(None)
3062    }
3063
3064    /// Whether this executor's backend can run the unified mixed prefill+decode
3065    /// forward natively. When false, the engine routes Qwen3-MoE batches through
3066    /// the legacy split path. Reported by the (backend-aware) executor so the
3067    /// engine stays backend-agnostic — replaces a `cfg(target_os)` branch that
3068    /// previously hard-coded "Metal/CPU lack native unified" in the hot path.
3069    ///
3070    /// Default false (conservative legacy path); accelerators with a native
3071    /// unified forward override to true.
3072    fn supports_native_unified_decode(&self) -> bool {
3073        false
3074    }
3075
3076    /// Per-request KV capacity in tokens when the executor owns a smaller
3077    /// runtime cache window than the model's declared context length.
3078    fn kv_capacity(&self) -> Option<usize> {
3079        None
3080    }
3081
3082    /// Device-aware limits selected from the exact compiled memory plan before
3083    /// static model upload. Legacy executors do not provide this evidence.
3084    fn startup_memory_plan(&self) -> Option<&ferrum_types::StartupMemoryPlan> {
3085        None
3086    }
3087
3088    /// Installs the product-owned execution event sink before requests start.
3089    ///
3090    /// Legacy executors have no typed execution journal and keep the default
3091    /// no-op. Executors backed by the vNext runtime retain this sink with each
3092    /// admitted request so node/operation events share the product artifact.
3093    fn attach_execution_event_sink(&self, _sink: Arc<dyn crate::vnext::ExecutionEventSink>) {}
3094
3095    /// Current plan-local capacity evidence for scheduler wake suppression.
3096    /// Legacy-engine executors return `None`; an executor declaring
3097    /// [`ExecutionResourceAuthority::PlanRuntime`] must return `Some`.
3098    fn execution_capacity_epochs(&self) -> Result<Option<ExecutorAdmissionEpochs>> {
3099        Ok(None)
3100    }
3101
3102    /// Writes the canonical per-source availability generations into
3103    /// caller-owned storage and returns the matching global audit epochs.
3104    /// Executors with typed dynamic admission override this to avoid allocating
3105    /// on steady scheduler ticks.
3106    fn write_execution_capacity_snapshot(
3107        &self,
3108        availability: &mut Vec<crate::vnext::CapacityAvailabilityEpoch>,
3109    ) -> Result<Option<ExecutorAdmissionEpochs>> {
3110        availability.clear();
3111        self.execution_capacity_epochs()
3112    }
3113
3114    /// Synchronously subscribes to every source named by one passive capacity
3115    /// wait. The returned registration must remain alive until it is awaited or
3116    /// deliberately cancelled by being dropped.
3117    ///
3118    /// Legacy-engine executors return `None`. An executor declaring
3119    /// [`ExecutionResourceAuthority::PlanRuntime`] must return `Some` for a
3120    /// wait condition issued by its own admission coordinator.
3121    fn register_execution_capacity_waiter(
3122        &self,
3123        _observed: &crate::vnext::CapacityWaitCondition,
3124    ) -> Result<Option<ExecutorCapacityWaitRegistration>> {
3125        Ok(None)
3126    }
3127
3128    /// Probe and retain the exact request/sequence authority needed by a
3129    /// future prefill. No provider encode, kernel launch, or device submit may
3130    /// occur in this method.
3131    fn try_admit_prefill(
3132        &self,
3133        _input: ExecutorPrefillAdmission<'_>,
3134    ) -> Result<ExecutorPrefillAdmissionDecision> {
3135        Err(ferrum_types::FerrumError::unsupported(
3136            "plan-runtime prefill admission is not implemented",
3137        ))
3138    }
3139
3140    /// Release an admitted but not yet active prefill authority.
3141    ///
3142    /// Returns true only when a retained authority was found and released.
3143    fn cancel_prefill_admission(&self, _request_id: &RequestId) -> bool {
3144        false
3145    }
3146
3147    /// Whether this exact plan can retain and restore independent prefix state.
3148    /// This must include resolved model/provider support and product policy.
3149    fn supports_plan_runtime_prefix_restore(&self) -> bool {
3150        false
3151    }
3152
3153    /// Restore a proper prefix into a freshly admitted request before the
3154    /// scheduler publishes any work for it. A miss submits no device work and
3155    /// preserves the admitted target. Successful restoration remains gated
3156    /// until the engine publishes the matching scheduler/executor progress and
3157    /// acknowledges the returned owner. Cancellation must retain unknown native
3158    /// writes until the existing completion path proves quiescence.
3159    async fn try_restore_plan_runtime_prefix(
3160        &self,
3161        _input: PlanRuntimePrefixRestoreInput<'_>,
3162    ) -> Result<PlanRuntimePrefixRestoreOutcome> {
3163        Ok(PlanRuntimePrefixRestoreOutcome::Unavailable)
3164    }
3165
3166    /// Writes the exact availability sources advanced when this request
3167    /// authority is preempted for recompute.
3168    ///
3169    /// `true` proves that `preemption` still identifies a live, quiescently
3170    /// releasable authority and that `sources` is its complete release
3171    /// footprint. Callers must treat `false` as not releasable; a generic
3172    /// "owns some cache" observation is not evidence that another request can
3173    /// advance the source on which the current frontier is blocked.
3174    fn write_execution_capacity_release_sources(
3175        &self,
3176        _preemption: &ExecutorExecutionCapacityPreemption,
3177        sources: &mut Vec<crate::vnext::CapacityAvailabilitySource>,
3178    ) -> Result<bool> {
3179        sources.clear();
3180        Ok(false)
3181    }
3182
3183    /// Retire one exact request-scoped runtime authority for recompute.
3184    ///
3185    /// Success is a terminal release fence: all provider/device work that can
3186    /// access the authority is quiescent and the request can be admitted as a
3187    /// new sequence incarnation. Implementations must fail closed on identity
3188    /// mismatch or an in-flight authority they cannot terminalize.
3189    async fn preempt_execution_capacity(
3190        &self,
3191        _preemption: ExecutorExecutionCapacityPreemption,
3192    ) -> Result<ExecutorExecutionCapacityPreemptionReceipt> {
3193        Err(FerrumError::unsupported(
3194            "request-scoped execution-capacity preemption is not implemented",
3195        ))
3196    }
3197
3198    /// Consume one retained logical/physical backing deferral after the
3199    /// scheduler waiting lock has been released. Implementations must perform
3200    /// at most one bounded maintenance attempt and publish capacity epochs only
3201    /// after real backing is installed.
3202    fn maintain_prefill_backing(
3203        &self,
3204        _request_id: &RequestId,
3205    ) -> Result<ExecutorPrefillMaintenanceOutcome> {
3206        Err(ferrum_types::FerrumError::unsupported(
3207            "plan-runtime prefill backing maintenance is not implemented",
3208        ))
3209    }
3210
3211    /// Reserve model-owned KV slots before a forward is dispatched.
3212    ///
3213    /// This is the executor-level admission hook for vLLM-style paged KV. The
3214    /// engine calls it at the batch boundary so a request that cannot grow its
3215    /// KV cache is delayed or preempted before kernel launch instead of
3216    /// panicking inside attention.
3217    fn reserve_kv_slots(&self, _requests: &[KvSlotRequest]) -> Result<Option<KvSlotReservation>> {
3218        Ok(None)
3219    }
3220
3221    /// Snapshot model-owned paged-KV capacity without allocating slots.
3222    ///
3223    /// Executors without model-owned paged KV return `None`.
3224    fn kv_slot_capacity_snapshot(&self) -> Option<KvSlotCapacitySnapshot> {
3225        None
3226    }
3227
3228    /// Recurrent-state allocation spec for this request, when the model has
3229    /// state-space or hybrid layers that need per-request recurrent state.
3230    ///
3231    /// Attention-only models return `None`. If this returns `Some`, the engine
3232    /// must allocate a recurrent-state handle before prefill and pass it through
3233    /// prefill/decode inputs. The default keeps existing executors KV-only.
3234    fn recurrent_state_spec(
3235        &self,
3236        _request_id: &RequestId,
3237        _input_tokens: &[TokenId],
3238    ) -> Result<Option<RecurrentStateSpec>> {
3239        Ok(None)
3240    }
3241
3242    /// Execute prefill phase (process initial prompt)
3243    async fn prefill(&self, input: &PrefillInput) -> Result<PrefillOutput>;
3244
3245    /// Execute one exact prefill chunk with an explicit pre-submit capacity
3246    /// deferral edge. Legacy executors inherit full-prefill behavior.
3247    async fn prefill_with_capacity(&self, input: &PrefillInput) -> Result<ExecutorPrefillOutcome> {
3248        let output = self.prefill(input).await?;
3249        let chunk = match input.chunk {
3250            Some(chunk) => chunk,
3251            None => PrefillChunk::new(0, input.sequence_length(), input.sequence_length())?,
3252        };
3253        Ok(ExecutorPrefillOutcome::Completed(
3254            ExecutorPrefillCompletion::exact(output, chunk),
3255        ))
3256    }
3257
3258    /// Batch prefill: process multiple prompts' prefill in ONE forward pass.
3259    ///
3260    /// Default implementation falls back to per-request `prefill()` (serial,
3261    /// which is the current behavior the engine sees today). Executors that
3262    /// support unified mixed-batch forward (e.g. via `model.unified_forward`
3263    /// over a varlen QKV path) should override this to amortize launch /
3264    /// kernel-overhead across all `inputs` items in one call.
3265    ///
3266    /// Used by the continuous-batching engine to coalesce a cohort of new
3267    /// prefills (apples M3 c=32 sees 32 simultaneous prefills as one logical
3268    /// batch; the serial fallback runs each in ~47 ms while a true batched
3269    /// path runs all 32 in ~100 ms).
3270    async fn batch_prefill(&self, inputs: &[PrefillInput]) -> Result<Vec<PrefillOutput>> {
3271        let mut outputs = Vec::with_capacity(inputs.len());
3272        for input in inputs {
3273            outputs.push(self.prefill(input).await?);
3274        }
3275        Ok(outputs)
3276    }
3277
3278    /// Attempt one physical, capacity-aware prefill batch.
3279    ///
3280    /// Implementations must either complete every input in original order or
3281    /// return `NotSubmitted` after restoring every retained prefill authority
3282    /// to a retryable state. Partial device submission is an ordinary error,
3283    /// never a `NotSubmitted` result.
3284    async fn batch_prefill_with_capacity(
3285        &self,
3286        _inputs: &[PrefillInput],
3287    ) -> Result<ExecutorBatchPrefillOutcome> {
3288        Ok(ExecutorBatchPrefillOutcome::Unsupported)
3289    }
3290
3291    /// Tensor-free prefill for executors that declare
3292    /// [`ExecutionResourceAuthority::PlanRuntime`].
3293    ///
3294    /// The default fails closed because adapting through [`PrefillInput`]
3295    /// would silently restore a host tensor boundary.
3296    async fn plan_runtime_prefill_with_capacity(
3297        &self,
3298        _input: &PlanRuntimePrefillInput,
3299    ) -> Result<PlanRuntimePrefillOutcome> {
3300        Err(FerrumError::unsupported(
3301            "tensor-free plan-runtime prefill is not implemented",
3302        ))
3303    }
3304
3305    /// Attempt one physical tensor-free prefill batch.
3306    ///
3307    /// `Unsupported` is an optimization fallback to the typed single-request
3308    /// method. `NotSubmitted` proves that no participant reached provider
3309    /// encode or device submission.
3310    async fn plan_runtime_batch_prefill_with_capacity(
3311        &self,
3312        _inputs: &[PlanRuntimePrefillInput],
3313    ) -> Result<PlanRuntimeBatchPrefillOutcome> {
3314        Ok(PlanRuntimeBatchPrefillOutcome::Unsupported)
3315    }
3316
3317    /// Discard an exact prefill authority after engine-side validation,
3318    /// sampling, or scheduler commit fails.
3319    ///
3320    /// Plan runtimes should override this to remove both retained-prefill and
3321    /// active state by the exact opaque handle, not by a reusable request id.
3322    fn discard_plan_runtime_prefill(&self, authority: PlanRuntimePrefillAuthority) -> Result<()> {
3323        self.release_cache(&authority.kv_cache().cache_id());
3324        Ok(())
3325    }
3326
3327    /// Execute decode phase (generate next token)
3328    async fn decode(&self, input: &DecodeInput) -> Result<DecodeOutput>;
3329
3330    /// Batch decode: process multiple sequences in one forward pass.
3331    ///
3332    /// A successful result must contain exactly one output per input, in the
3333    /// original input order, and each output cache must retain the identity of
3334    /// its corresponding input cache. Implementations must not expose partial
3335    /// success as a shorter or reordered vector.
3336    ///
3337    /// The default implementation falls back to serial per-request `decode()`.
3338    /// Executors with a typed batch submission path should override this so one
3339    /// call maps to one resource step and one terminal submission fence.
3340    async fn batch_decode(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>> {
3341        let mut outputs = Vec::with_capacity(inputs.len());
3342        for input in inputs {
3343            outputs.push(self.decode(input).await?);
3344        }
3345        Ok(outputs)
3346    }
3347
3348    /// Batch decode with an explicit pre-submit capacity deferral edge.
3349    ///
3350    /// Legacy executors inherit the successful/error-only behavior. A runtime
3351    /// with typed resource authority overrides this method so temporary
3352    /// capacity pressure is never flattened into a stringly resource error.
3353    async fn batch_decode_with_capacity(
3354        &self,
3355        inputs: &[DecodeInput],
3356    ) -> Result<ExecutorBatchDecodeOutcome> {
3357        self.batch_decode(inputs)
3358            .await
3359            .map(ExecutorBatchDecodeOutcome::Completed)
3360    }
3361
3362    /// Tensor-free batch decode for executors that declare
3363    /// [`ExecutionResourceAuthority::PlanRuntime`].
3364    ///
3365    /// Implementations must preserve input ordering and cache identity exactly.
3366    /// Temporary pressure may return `Deferred` only before device submission.
3367    /// The default fails closed because adapting through [`DecodeInput`] would
3368    /// silently restore the host-tensor boundary this contract removes.
3369    async fn plan_runtime_batch_decode_with_capacity(
3370        &self,
3371        _inputs: &[PlanRuntimeDecodeInput],
3372    ) -> Result<PlanRuntimeBatchDecodeOutcome> {
3373        Err(FerrumError::unsupported(
3374            "tensor-free plan-runtime batch decode is not implemented",
3375        ))
3376    }
3377
3378    /// Unified mixed-batch forward: process a [`UnifiedBatch`] containing
3379    /// any combination of prefill chunks (one or more `q_tokens` per item,
3380    /// possibly continuing from `pos_offset > 0`) and decode steps
3381    /// (`q_tokens.len() == 1`, `is_final_chunk = true`) in a single model
3382    /// forward pass.
3383    ///
3384    /// Returns one element per `batch.items[i]`:
3385    /// - `Some(logits)` for items with `is_final_chunk = true` (the
3386    ///   request's final-position logits, ready for sampling)
3387    /// - `None` for intermediate prefill chunks (no lm_head executed —
3388    ///   model only updates KV state)
3389    ///
3390    /// Default implementation returns `Err(unsupported)`. Concrete LLM
3391    /// executors should override with either:
3392    /// - A behavioral fallback that dispatches each chunk via existing
3393    ///   `prefill()` and groups decode items into `batch_decode()` (this
3394    ///   preserves current behavior; no perf change), OR
3395    /// - A real unified-forward path that runs all items through one
3396    ///   `[M_total, hidden]` GEMM chain with a varlen attention kernel
3397    ///   (this is the chunked-prefill perf unlock).
3398    async fn unified_decode(&self, _batch: &UnifiedBatch) -> Result<Vec<Option<Vec<f32>>>> {
3399        Err(ferrum_types::FerrumError::unsupported(
3400            "unified_decode not implemented for this executor",
3401        ))
3402    }
3403
3404    /// Optional: full forward pass (for non-autoregressive use cases)
3405    async fn forward(&self, _input: &TensorRef) -> Result<TensorRef> {
3406        // Default implementation not supported
3407        Err(ferrum_types::FerrumError::unsupported(
3408            "Full forward pass not supported by this executor",
3409        ))
3410    }
3411
3412    /// Roll the KV cache for this executor's sequence back to `new_len`.
3413    /// Used by speculative decoding on partial rejection so the next
3414    /// iteration sees a KV prefix that matches the accepted token stream.
3415    /// Default: Ok(()) — executors that don't cache per-sequence state
3416    /// (stub, mock) are inherently tolerant; real LLM executors override.
3417    async fn truncate_kv(
3418        &self,
3419        _kv_cache: &std::sync::Arc<dyn crate::KvCacheHandle>,
3420        _new_len: usize,
3421    ) -> Result<()> {
3422        Ok(())
3423    }
3424
3425    /// Multi-position decode-verify: one forward over `N+1` tokens,
3426    /// producing one logits row per position. Used by speculative
3427    /// decoding's target path so we don't pay N+1 sequential forwards.
3428    ///
3429    /// Default falls back to N+1 sequential `decode()` calls — correct
3430    /// but slow; real LLM executors override.
3431    ///
3432    /// Returns a `Vec<DecodeOutput>` of length `inputs.len()` with the
3433    /// final KV handle attached to the last element.
3434    async fn forward_verify(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>> {
3435        let mut out = Vec::with_capacity(inputs.len());
3436        for input in inputs {
3437            out.push(self.decode(input).await?);
3438        }
3439        Ok(out)
3440    }
3441
3442    /// Get executor capabilities
3443    fn capabilities(&self) -> ExecutorCapabilities;
3444
3445    /// Get current executor status
3446    fn status(&self) -> ExecutorStatus;
3447
3448    /// Optional model/executor cache metrics.
3449    ///
3450    /// Concrete LLM executors use this for model-level paged KV prefix reuse
3451    /// counters. Default implementations keep non-autoregressive executors
3452    /// and tests from needing cache-specific plumbing.
3453    fn cache_metrics_snapshot(&self) -> Option<serde_json::Value> {
3454        None
3455    }
3456
3457    /// Optional compact provider-attribution witness emitted by executors
3458    /// whose immutable plan can bind quantized source tensors to selected
3459    /// operation providers without exposing per-tensor logs.
3460    fn execution_attribution_snapshot(&self) -> Option<serde_json::Value> {
3461        None
3462    }
3463
3464    /// Optional LoRA runtime metrics.
3465    fn lora_metrics_snapshot(&self) -> Option<serde_json::Value> {
3466        None
3467    }
3468
3469    /// Complete executor-owned startup preparation before either product
3470    /// entrypoint can accept a request.
3471    ///
3472    /// Implementations use this cold-path hook for work that requires the
3473    /// fully constructed executor but must not be charged to a user's first
3474    /// request, such as compiling reusable execution shapes. The default is a
3475    /// no-op so existing executors remain source compatible.
3476    async fn prepare_startup(&self) -> Result<()> {
3477        Ok(())
3478    }
3479
3480    /// Warm up executor (load model, allocate memory, etc.)
3481    async fn warmup(&mut self) -> Result<()> {
3482        // Default no-op implementation
3483        Ok(())
3484    }
3485
3486    /// Shutdown executor gracefully
3487    async fn shutdown(&mut self) -> Result<()> {
3488        // Default no-op implementation
3489        Ok(())
3490    }
3491
3492    /// Complete and release one cache authority with product-authoritative
3493    /// terminal token counts.
3494    ///
3495    /// Legacy executors only need physical release and inherit that behavior.
3496    /// Plan runtimes with terminal journals override this method so completion
3497    /// cannot be inferred from a generic release operation. Implementations may
3498    /// await retention of the last completed native state before releasing its
3499    /// source; sampled output alone is not evidence of an executed frontier.
3500    async fn complete_cache(&self, completion: ExecutorSequenceCompletion) -> Result<()> {
3501        self.release_cache(completion.cache_id());
3502        Ok(())
3503    }
3504
3505    /// Release KV cache and state without asserting successful completion.
3506    ///
3507    /// Called for cancellation, failure, recompute, and legacy cleanup. The
3508    /// `cache_id` matches the value embedded in the `KvCacheHandle` returned by
3509    /// prefill/decode. Successful product completion uses [`Self::complete_cache`].
3510    fn release_cache(&self, _cache_id: &str) {
3511        // Default no-op — executors that manage per-sequence KV caches should override.
3512    }
3513}
3514
3515/// Executor capabilities and configuration
3516#[derive(Debug, Clone, Serialize, Deserialize)]
3517pub struct ExecutorCapabilities {
3518    /// Maximum supported batch size
3519    pub max_batch_size: usize,
3520    /// Maximum sequence length
3521    pub max_sequence_length: usize,
3522    /// Supported attention mechanisms
3523    pub attention_mechanisms: Vec<AttentionType>,
3524    /// Whether executor supports dynamic batching
3525    pub supports_dynamic_batching: bool,
3526    /// Whether executor supports continuous batching
3527    pub supports_continuous_batching: bool,
3528    /// Whether executor supports speculative decoding
3529    pub supports_speculative_decoding: bool,
3530    /// Whether executor supports tensor parallelism
3531    pub supports_tensor_parallelism: bool,
3532    /// Whether executor supports pipeline parallelism
3533    pub supports_pipeline_parallelism: bool,
3534    /// Supported data types
3535    pub supported_dtypes: Vec<ferrum_types::DataType>,
3536    /// Supported devices
3537    pub supported_devices: Vec<ferrum_types::Device>,
3538    /// Memory requirements estimation
3539    pub memory_requirements: MemoryRequirements,
3540}
3541
3542/// Attention mechanism types
3543#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
3544pub enum AttentionType {
3545    /// Standard multi-head attention
3546    MultiHead,
3547    /// Multi-query attention (MQA)
3548    MultiQuery,
3549    /// Grouped-query attention (GQA)
3550    GroupedQuery,
3551    /// Flash attention
3552    Flash,
3553    /// Paged attention
3554    Paged,
3555    /// Sliding window attention
3556    SlidingWindow,
3557}
3558
3559/// Logical state bytes across the complete model, before physical page
3560/// alignment, pool residency, checkpoint copies or operation workspace.
3561#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
3562pub struct TypedSequenceStateMemory {
3563    /// K/V payload and quantization scale bytes for one token in one sequence.
3564    pub kv_bytes_per_token: u64,
3565    /// Other sequence states whose size scales with token count.
3566    pub other_token_scaled_bytes_per_token: u64,
3567    /// Fixed sequence state, including recurrent accumulators and windows.
3568    pub fixed_bytes_per_sequence: u64,
3569}
3570
3571/// Memory requirements for model execution. This describes logical model
3572/// requirements, not measured device allocation or an admission guarantee.
3573#[derive(Debug, Clone, Deserialize)]
3574pub struct MemoryRequirements {
3575    /// Model parameter memory in bytes
3576    pub parameter_memory: u64,
3577    /// Minimum activation memory per token
3578    pub activation_memory_per_token: usize,
3579    /// Legacy KV cache bytes per token per layer. Ignored and omitted from
3580    /// serialization when typed_sequence_state supplies complete-model bytes.
3581    #[serde(default)]
3582    pub kv_cache_memory_per_token: usize,
3583    /// Exact logical sequence state from a typed model plan. Old wire values
3584    /// without this field retain the legacy per-layer interpretation.
3585    #[serde(default)]
3586    pub typed_sequence_state: Option<TypedSequenceStateMemory>,
3587    /// Additional overhead memory
3588    pub overhead_memory: u64,
3589}
3590
3591impl Serialize for MemoryRequirements {
3592    fn serialize<S: serde::Serializer>(
3593        &self,
3594        serializer: S,
3595    ) -> std::result::Result<S::Ok, S::Error> {
3596        use serde::ser::SerializeStruct;
3597        let mut value = serializer.serialize_struct("MemoryRequirements", 4)?;
3598        value.serialize_field("parameter_memory", &self.parameter_memory)?;
3599        value.serialize_field(
3600            "activation_memory_per_token",
3601            &self.activation_memory_per_token,
3602        )?;
3603        match &self.typed_sequence_state {
3604            Some(state) => value.serialize_field("typed_sequence_state", state)?,
3605            None => value
3606                .serialize_field("kv_cache_memory_per_token", &self.kv_cache_memory_per_token)?,
3607        }
3608        value.serialize_field("overhead_memory", &self.overhead_memory)?;
3609        value.end()
3610    }
3611}
3612
3613impl MemoryRequirements {
3614    /// Calculate a logical memory estimate, returning None on arithmetic
3615    /// overflow. Typed sequence state is already summed across all layers.
3616    pub fn checked_calculate_total_memory(
3617        &self,
3618        batch_size: usize,
3619        sequence_length: usize,
3620        num_layers: usize,
3621    ) -> Option<u64> {
3622        let batch = u64::try_from(batch_size).ok()?;
3623        let tokens = u64::try_from(sequence_length).ok()?;
3624        let token_count = batch.checked_mul(tokens)?;
3625        let activation_mem = u64::try_from(self.activation_memory_per_token)
3626            .ok()?
3627            .checked_mul(token_count)?;
3628        let state_mem = match self.typed_sequence_state {
3629            Some(state) => state
3630                .kv_bytes_per_token
3631                .checked_add(state.other_token_scaled_bytes_per_token)?
3632                .checked_mul(token_count)?
3633                .checked_add(state.fixed_bytes_per_sequence.checked_mul(batch)?)?,
3634            None => u64::try_from(self.kv_cache_memory_per_token)
3635                .ok()?
3636                .checked_mul(token_count)?
3637                .checked_mul(u64::try_from(num_layers).ok()?)?,
3638        };
3639        self.parameter_memory
3640            .checked_add(activation_mem)?
3641            .checked_add(state_mem)?
3642            .checked_add(self.overhead_memory)
3643    }
3644
3645    /// Calculate a logical memory estimate. Overflow saturates to u64::MAX;
3646    /// callers needing an explicit overflow result can use the checked method.
3647    pub fn calculate_total_memory(
3648        &self,
3649        batch_size: usize,
3650        sequence_length: usize,
3651        num_layers: usize,
3652    ) -> u64 {
3653        self.checked_calculate_total_memory(batch_size, sequence_length, num_layers)
3654            .unwrap_or(u64::MAX)
3655    }
3656}
3657
3658#[cfg(test)]
3659mod memory_requirements_tests {
3660    use super::{MemoryRequirements, TypedSequenceStateMemory};
3661
3662    #[test]
3663    fn typed_sequence_memory_counts_scales_and_fixed_state_once_per_sequence() {
3664        let memory = MemoryRequirements {
3665            parameter_memory: 100,
3666            activation_memory_per_token: 4,
3667            kv_cache_memory_per_token: 999,
3668            typed_sequence_state: Some(TypedSequenceStateMemory {
3669                kv_bytes_per_token: 528,
3670                other_token_scaled_bytes_per_token: 8,
3671                fixed_bytes_per_sequence: 64,
3672            }),
3673            overhead_memory: 20,
3674        };
3675        let expected = 100 + 4 * 2 * 3 + (528 + 8) * 2 * 3 + 64 * 2 + 20;
3676        for layers in [1, 3, 32] {
3677            assert_eq!(
3678                memory.checked_calculate_total_memory(2, 3, layers),
3679                Some(expected)
3680            );
3681        }
3682        let wire = serde_json::to_value(&memory).unwrap();
3683        assert!(wire.get("kv_cache_memory_per_token").is_none());
3684        assert_eq!(wire["typed_sequence_state"]["kv_bytes_per_token"], 528);
3685        let decoded: MemoryRequirements = serde_json::from_value(wire).unwrap();
3686        assert_eq!(decoded.calculate_total_memory(2, 3, 32), expected);
3687    }
3688
3689    #[test]
3690    fn legacy_memory_wire_keeps_per_layer_calculation_and_rejects_overflow() {
3691        let wire = serde_json::json!({
3692            "parameter_memory": 100,
3693            "activation_memory_per_token": 4,
3694            "kv_cache_memory_per_token": 16,
3695            "overhead_memory": 20,
3696        });
3697        let mut memory: MemoryRequirements = serde_json::from_value(wire.clone()).unwrap();
3698        assert!(memory.typed_sequence_state.is_none());
3699        assert_eq!(
3700            memory.calculate_total_memory(2, 3, 5),
3701            100 + 4 * 2 * 3 + 16 * 2 * 3 * 5 + 20
3702        );
3703        assert_eq!(serde_json::to_value(&memory).unwrap(), wire);
3704        memory.parameter_memory = u64::MAX;
3705        assert_eq!(memory.checked_calculate_total_memory(1, 1, 1), None);
3706        assert_eq!(memory.calculate_total_memory(1, 1, 1), u64::MAX);
3707        memory.parameter_memory = 0;
3708        memory.typed_sequence_state = Some(TypedSequenceStateMemory {
3709            kv_bytes_per_token: u64::MAX,
3710            other_token_scaled_bytes_per_token: 1,
3711            fixed_bytes_per_sequence: 0,
3712        });
3713        assert_eq!(memory.checked_calculate_total_memory(1, 1, 1), None);
3714    }
3715}
3716
3717/// Executor status information
3718#[derive(Debug, Clone, Serialize, Deserialize)]
3719pub struct ExecutorStatus {
3720    /// Current executor state
3721    pub state: ExecutorState,
3722    /// Whether executor is ready to accept requests
3723    pub is_ready: bool,
3724    /// Current batch size being processed
3725    pub current_batch_size: usize,
3726    /// Number of prefill operations completed
3727    pub prefill_operations: u64,
3728    /// Number of decode operations completed
3729    pub decode_operations: u64,
3730    /// Average prefill time in milliseconds
3731    pub avg_prefill_time_ms: f64,
3732    /// Average decode time in milliseconds
3733    pub avg_decode_time_ms: f64,
3734    /// Memory usage statistics
3735    pub memory_usage: ExecutorMemoryUsage,
3736    /// Last operation timestamp
3737    #[serde(skip)]
3738    pub last_operation: Option<std::time::Instant>,
3739}
3740
3741/// Executor state
3742#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
3743pub enum ExecutorState {
3744    /// Executor is initializing
3745    Initializing,
3746    /// Executor is ready to accept requests
3747    Ready,
3748    /// Executor is processing requests
3749    Busy,
3750    /// Executor encountered an error
3751    Error,
3752    /// Executor is shutting down
3753    Shutdown,
3754}
3755
3756/// Executor memory usage
3757#[derive(Debug, Clone, Serialize, Deserialize)]
3758pub struct ExecutorMemoryUsage {
3759    /// Total allocated memory in bytes
3760    pub allocated_bytes: usize,
3761    /// Currently used memory in bytes
3762    pub used_bytes: usize,
3763    /// Peak memory usage
3764    pub peak_bytes: usize,
3765    /// Memory utilization percentage
3766    pub utilization_percent: f32,
3767}
3768
3769/// Batch model executor for processing multiple requests efficiently
3770#[async_trait]
3771pub trait BatchModelExecutor: ModelExecutor {
3772    /// Execute batch prefill for multiple sequences
3773    async fn batch_prefill(&self, inputs: &[PrefillInput]) -> Result<Vec<PrefillOutput>>;
3774
3775    /// Execute batch decode for multiple sequences
3776    async fn batch_decode(&self, inputs: &[DecodeInput]) -> Result<Vec<DecodeOutput>>;
3777
3778    /// Get optimal batch size for current conditions
3779    fn optimal_batch_size(&self) -> usize;
3780
3781    /// Check if batch size is supported
3782    fn supports_batch_size(&self, batch_size: usize) -> bool;
3783}
3784
3785/// Speculative execution support
3786#[async_trait]
3787pub trait SpeculativeExecutor: ModelExecutor {
3788    /// Execute speculative decoding with draft model
3789    async fn speculative_decode(
3790        &self,
3791        input: &DecodeInput,
3792        draft_tokens: &[ferrum_types::TokenId],
3793        acceptance_threshold: f32,
3794    ) -> Result<SpeculativeDecodeOutput>;
3795}
3796
3797/// Output from speculative decoding
3798#[derive(Debug, Clone)]
3799pub struct SpeculativeDecodeOutput {
3800    /// Accepted tokens (subset of draft tokens)
3801    pub accepted_tokens: Vec<ferrum_types::TokenId>,
3802    /// Logits for the next token after last accepted
3803    pub next_logits: TensorRef,
3804    /// Updated KV cache
3805    pub kv_cache: Arc<dyn KvCacheHandle>,
3806    /// Number of draft tokens accepted
3807    pub acceptance_count: usize,
3808}
3809
3810/// Model executor factory
3811#[async_trait]
3812pub trait ModelExecutorFactory: Send + Sync {
3813    /// Create executor from model configuration
3814    async fn create_executor(&self, config: &ExecutorConfig) -> Result<Box<dyn ModelExecutor>>;
3815
3816    /// Create batch executor
3817    async fn create_batch_executor(
3818        &self,
3819        config: &ExecutorConfig,
3820    ) -> Result<Box<dyn BatchModelExecutor>>;
3821
3822    /// Get supported executor types
3823    fn supported_types(&self) -> Vec<ExecutorType>;
3824
3825    /// Validate configuration
3826    fn validate_config(&self, config: &ExecutorConfig) -> Result<()>;
3827}
3828
3829/// Executor configuration
3830#[derive(Debug, Clone, Serialize, Deserialize)]
3831pub struct ExecutorConfig {
3832    /// Model information
3833    pub model_info: ModelInfo,
3834    /// Target device
3835    pub device: ferrum_types::Device,
3836    /// Data type for computation
3837    pub dtype: ferrum_types::DataType,
3838    /// Maximum batch size
3839    pub max_batch_size: usize,
3840    /// Maximum sequence length
3841    pub max_sequence_length: usize,
3842    /// Attention configuration
3843    pub attention_config: ExecutorAttentionConfig,
3844    /// Memory configuration
3845    pub memory_config: ExecutorMemoryConfig,
3846    /// Optimization settings
3847    pub optimization_config: OptimizationConfig,
3848    /// Additional executor-specific options
3849    pub executor_options: HashMap<String, serde_json::Value>,
3850}
3851
3852/// Runtime attention configuration for model executor
3853///
3854/// Note: This is different from ferrum_types::AttentionConfig which describes
3855/// the model architecture's attention configuration from config.json.
3856/// This type describes the runtime execution settings.
3857#[derive(Debug, Clone, Serialize, Deserialize)]
3858pub struct ExecutorAttentionConfig {
3859    /// Type of attention to use
3860    pub attention_type: AttentionType,
3861    /// Enable flash attention if available
3862    pub enable_flash_attention: bool,
3863    /// Enable paged attention
3864    pub enable_paged_attention: bool,
3865    /// Block size for paged attention
3866    pub block_size: Option<usize>,
3867    /// Sliding window size (if using sliding window attention)
3868    pub sliding_window_size: Option<usize>,
3869}
3870
3871/// Memory configuration for executor
3872#[derive(Debug, Clone, Serialize, Deserialize)]
3873pub struct ExecutorMemoryConfig {
3874    /// Enable memory pooling
3875    pub enable_memory_pooling: bool,
3876    /// Memory pool size in bytes (None for auto)
3877    pub memory_pool_size: Option<usize>,
3878    /// Enable KV cache sharing
3879    pub enable_kv_cache_sharing: bool,
3880    /// Maximum memory usage percentage
3881    pub max_memory_usage: f32,
3882}
3883
3884/// Optimization configuration
3885#[derive(Debug, Clone, Serialize, Deserialize)]
3886pub struct OptimizationConfig {
3887    /// Enable CUDA graphs (if supported)
3888    pub enable_cuda_graphs: bool,
3889    /// Enable kernel fusion
3890    pub enable_kernel_fusion: bool,
3891    /// Enable mixed precision
3892    pub enable_mixed_precision: bool,
3893    /// Optimization level (0-3)
3894    pub optimization_level: u8,
3895    /// Custom optimization flags
3896    pub custom_flags: HashMap<String, bool>,
3897}
3898
3899/// Supported executor types
3900#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
3901pub enum ExecutorType {
3902    /// Standard sequential executor
3903    Sequential,
3904    /// Batch executor for parallel processing
3905    Batch,
3906    /// Continuous batching executor
3907    ContinuousBatch,
3908    /// Speculative decoding executor
3909    Speculative,
3910    /// Pipeline parallel executor
3911    PipelineParallel,
3912    /// Tensor parallel executor
3913    TensorParallel,
3914}
3915
3916/// Executor performance metrics
3917#[derive(Debug, Clone, Serialize, Deserialize)]
3918pub struct ExecutorMetrics {
3919    /// Total operations executed
3920    pub total_operations: u64,
3921    /// Prefill operations
3922    pub prefill_operations: u64,
3923    /// Decode operations
3924    pub decode_operations: u64,
3925    /// Average prefill latency (ms)
3926    pub avg_prefill_latency: f64,
3927    /// Average decode latency (ms)
3928    pub avg_decode_latency: f64,
3929    /// P95 prefill latency (ms)
3930    pub p95_prefill_latency: f64,
3931    /// P95 decode latency (ms)
3932    pub p95_decode_latency: f64,
3933    /// Throughput (tokens per second)
3934    pub throughput_tps: f64,
3935    /// Memory efficiency (used/allocated)
3936    pub memory_efficiency: f32,
3937    /// Batch utilization
3938    pub batch_utilization: f32,
3939}
3940
3941/// Executor registry for managing multiple executors
3942pub trait ExecutorRegistry: Send + Sync {
3943    /// Register executor with name
3944    fn register(&mut self, name: &str, executor: Box<dyn ModelExecutor>) -> Result<()>;
3945
3946    /// Get executor by name
3947    fn get(&self, name: &str) -> Option<&dyn ModelExecutor>;
3948
3949    /// Remove executor by name
3950    fn remove(&mut self, name: &str) -> Option<Box<dyn ModelExecutor>>;
3951
3952    /// List registered executor names
3953    fn list_names(&self) -> Vec<String>;
3954
3955    /// Get executor metrics
3956    fn get_metrics(&self, name: &str) -> Option<ExecutorMetrics>;
3957}