1use std::{collections::BTreeSet, num::NonZeroUsize};
4
5use eredu_core::{
6 SpeculativeExecutionTopology, SpeculativeLifecycleObserver, SpeculativeLifecycleStage,
7 SpeculativeOutputError, TokenizerCompatibilityProof,
8};
9
10#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
12pub struct SpeculativeIdentity(String);
13
14impl SpeculativeIdentity {
15 pub fn new(value: impl Into<String>) -> Result<Self, SpeculativeContractError> {
17 let value = value.into();
18 if value.trim().is_empty() {
19 return Err(SpeculativeContractError::new(
20 "speculative identity must not be empty",
21 ));
22 }
23 Ok(Self(value))
24 }
25
26 pub fn as_str(&self) -> &str {
28 &self.0
29 }
30}
31
32#[derive(Debug, Clone, Eq, PartialEq)]
34pub struct SpeculativeCaptureEntry {
35 path: SpeculativeIdentity,
36 shape: Vec<usize>,
37 bounded_dimensions: BTreeSet<usize>,
38 owner: SpeculativeIdentity,
39 observation: SpeculativeIdentity,
40}
41
42impl SpeculativeCaptureEntry {
43 pub fn new(
45 path: SpeculativeIdentity,
46 shape: Vec<usize>,
47 owner: SpeculativeIdentity,
48 observation: SpeculativeIdentity,
49 ) -> Result<Self, SpeculativeContractError> {
50 if shape.is_empty() || shape.contains(&0) {
51 return Err(SpeculativeContractError::new(
52 "speculative capture shape must have positive extents",
53 ));
54 }
55 Ok(Self {
56 path,
57 shape,
58 bounded_dimensions: BTreeSet::new(),
59 owner,
60 observation,
61 })
62 }
63
64 pub fn with_bounded_dimension(
66 mut self,
67 dimension: usize,
68 ) -> Result<Self, SpeculativeContractError> {
69 if dimension >= self.shape.len() {
70 return Err(SpeculativeContractError::new(
71 "bounded capture dimension is outside the declared shape",
72 ));
73 }
74 self.bounded_dimensions.insert(dimension);
75 Ok(self)
76 }
77
78 pub const fn path(&self) -> &SpeculativeIdentity {
80 &self.path
81 }
82
83 pub fn shape(&self) -> &[usize] {
85 &self.shape
86 }
87
88 pub const fn bounded_dimensions(&self) -> &BTreeSet<usize> {
90 &self.bounded_dimensions
91 }
92
93 pub const fn owner(&self) -> &SpeculativeIdentity {
95 &self.owner
96 }
97
98 pub const fn observation(&self) -> &SpeculativeIdentity {
100 &self.observation
101 }
102}
103
104#[derive(Debug, Clone, Eq, PartialEq)]
106pub struct SpeculativeCaptureSchema {
107 identity: SpeculativeIdentity,
108 entries: Vec<SpeculativeCaptureEntry>,
109}
110
111impl SpeculativeCaptureSchema {
112 pub fn new(
114 identity: SpeculativeIdentity,
115 entries: impl IntoIterator<Item = SpeculativeCaptureEntry>,
116 ) -> Result<Self, SpeculativeContractError> {
117 let entries = entries.into_iter().collect::<Vec<_>>();
118 if entries.is_empty() {
119 return Err(SpeculativeContractError::new(
120 "speculative capture schema must not be empty",
121 ));
122 }
123 let paths = entries
124 .iter()
125 .map(SpeculativeCaptureEntry::path)
126 .collect::<BTreeSet<_>>();
127 if paths.len() != entries.len() {
128 return Err(SpeculativeContractError::new(
129 "speculative capture paths must be unique",
130 ));
131 }
132 let observations = entries
133 .iter()
134 .map(SpeculativeCaptureEntry::observation)
135 .collect::<BTreeSet<_>>();
136 if observations.len() != entries.len() {
137 return Err(SpeculativeContractError::new(
138 "speculative capture observation identities must be unique",
139 ));
140 }
141 Ok(Self { identity, entries })
142 }
143
144 pub const fn identity(&self) -> &SpeculativeIdentity {
146 &self.identity
147 }
148
149 pub fn entries(&self) -> &[SpeculativeCaptureEntry] {
151 &self.entries
152 }
153
154 pub fn instantiate(
156 &self,
157 shapes: impl IntoIterator<Item = Vec<usize>>,
158 ) -> Result<Self, SpeculativeCaptureError> {
159 let shapes = shapes.into_iter().collect::<Vec<_>>();
160 if shapes.len() != self.entries.len() {
161 return Err(SpeculativeCaptureError::ValueCount {
162 expected: self.entries.len(),
163 actual: shapes.len(),
164 });
165 }
166 let mut entries = Vec::with_capacity(self.entries.len());
167 for (expected, shape) in self.entries.iter().zip(shapes) {
168 if shape.len() != expected.shape.len()
169 || shape.contains(&0)
170 || shape.iter().enumerate().any(|(dimension, actual)| {
171 let maximum = expected.shape[dimension];
172 if expected.bounded_dimensions.contains(&dimension) {
173 *actual > maximum
174 } else {
175 *actual != maximum
176 }
177 })
178 {
179 return Err(SpeculativeCaptureError::ShapeMismatch);
180 }
181 let mut entry = expected.clone();
182 entry.shape = shape;
183 entry.bounded_dimensions.clear();
184 entries.push(entry);
185 }
186 Ok(Self {
187 identity: self.identity.clone(),
188 entries,
189 })
190 }
191
192 fn admits(&self, actual: &Self) -> bool {
193 self.identity == actual.identity
194 && self.entries.len() == actual.entries.len()
195 && self
196 .entries
197 .iter()
198 .zip(&actual.entries)
199 .all(|(expected, actual)| {
200 expected.path == actual.path
201 && expected.owner == actual.owner
202 && expected.observation == actual.observation
203 && expected.shape.len() == actual.shape.len()
204 && actual.bounded_dimensions.is_empty()
205 && expected
206 .shape
207 .iter()
208 .enumerate()
209 .all(|(dimension, maximum)| {
210 let actual = actual.shape[dimension];
211 actual > 0
212 && if expected.bounded_dimensions.contains(&dimension) {
213 actual <= *maximum
214 } else {
215 actual == *maximum
216 }
217 })
218 })
219 }
220}
221
222#[derive(Debug, Clone, Eq, PartialEq)]
224pub struct SpeculativeCaptureMetadata {
225 schema: SpeculativeCaptureSchema,
226 generation: u64,
227}
228
229impl SpeculativeCaptureMetadata {
230 pub const fn new(schema: SpeculativeCaptureSchema, generation: u64) -> Self {
232 Self { schema, generation }
233 }
234
235 pub const fn schema(&self) -> &SpeculativeCaptureSchema {
237 &self.schema
238 }
239
240 pub const fn generation(&self) -> u64 {
242 self.generation
243 }
244}
245
246#[derive(Debug, Clone, Eq, PartialEq)]
248pub struct SpeculativeCaptureEnvelope<T> {
249 metadata: SpeculativeCaptureMetadata,
250 values: Vec<T>,
251}
252
253impl<T> SpeculativeCaptureEnvelope<T> {
254 pub fn new(
256 metadata: SpeculativeCaptureMetadata,
257 values: Vec<T>,
258 ) -> Result<Self, SpeculativeCaptureError> {
259 if values.len() != metadata.schema.entries.len() {
260 return Err(SpeculativeCaptureError::ValueCount {
261 expected: metadata.schema.entries.len(),
262 actual: values.len(),
263 });
264 }
265 Ok(Self { metadata, values })
266 }
267
268 pub const fn metadata(&self) -> &SpeculativeCaptureMetadata {
270 &self.metadata
271 }
272
273 pub fn values(&self) -> &[T] {
275 &self.values
276 }
277
278 pub fn into_values(self) -> Vec<T> {
280 self.values
281 }
282
283 pub fn validate_against(
285 &self,
286 schema: &SpeculativeCaptureSchema,
287 generation: u64,
288 ) -> Result<(), SpeculativeCaptureError> {
289 if !schema.admits(&self.metadata.schema) {
290 return Err(SpeculativeCaptureError::SchemaMismatch);
291 }
292 if self.metadata.generation != generation {
293 return Err(SpeculativeCaptureError::GenerationMismatch {
294 expected: generation,
295 actual: self.metadata.generation,
296 });
297 }
298 Ok(())
299 }
300}
301
302#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
304pub enum SpeculativeCaptureError {
305 #[error("speculative capture expected {expected} values, received {actual}")]
307 ValueCount {
308 expected: usize,
310 actual: usize,
312 },
313 #[error("speculative capture schema does not match the selected architecture")]
315 SchemaMismatch,
316 #[error("speculative capture shape does not match the selected architecture")]
318 ShapeMismatch,
319 #[error("speculative capture generation {actual} does not match {expected}")]
321 GenerationMismatch {
322 expected: u64,
324 actual: u64,
326 },
327 #[error("speculative lane identity does not match the selected realization")]
329 RealizationMismatch,
330}
331
332#[derive(Debug, Clone, Copy, Eq, PartialEq)]
334#[non_exhaustive]
335pub enum SpeculativeStrategyClass {
336 EmbeddedSequential,
338 EmbeddedFused,
340 External,
342}
343
344#[derive(Debug, Clone, Eq, PartialEq)]
346pub struct SpeculativeStrategyRequirements {
347 class: SpeculativeStrategyClass,
348 identity: SpeculativeIdentity,
349 proposal_capacity: NonZeroUsize,
350 tokenizer_fingerprint: Option<[u8; 32]>,
351}
352
353impl SpeculativeStrategyRequirements {
354 pub fn embedded(
356 class: SpeculativeStrategyClass,
357 identity: SpeculativeIdentity,
358 proposal_capacity: NonZeroUsize,
359 ) -> Result<Self, SpeculativeContractError> {
360 match class {
361 SpeculativeStrategyClass::EmbeddedSequential
362 | SpeculativeStrategyClass::EmbeddedFused => Ok(Self {
363 class,
364 identity,
365 proposal_capacity,
366 tokenizer_fingerprint: None,
367 }),
368 SpeculativeStrategyClass::External => Err(SpeculativeContractError::new(
369 "external strategy requires an external-tokenizer identity",
370 )),
371 }
372 }
373
374 pub const fn external(
376 identity: SpeculativeIdentity,
377 proposal_capacity: NonZeroUsize,
378 tokenizer_fingerprint: [u8; 32],
379 ) -> Self {
380 Self {
381 class: SpeculativeStrategyClass::External,
382 identity,
383 proposal_capacity,
384 tokenizer_fingerprint: Some(tokenizer_fingerprint),
385 }
386 }
387
388 pub const fn class(&self) -> SpeculativeStrategyClass {
390 self.class
391 }
392
393 pub const fn identity(&self) -> &SpeculativeIdentity {
395 &self.identity
396 }
397
398 pub const fn proposal_capacity(&self) -> NonZeroUsize {
400 self.proposal_capacity
401 }
402
403 pub const fn tokenizer_fingerprint(&self) -> Option<[u8; 32]> {
405 self.tokenizer_fingerprint
406 }
407}
408
409#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
411#[non_exhaustive]
412pub enum SpeculativeMechanism {
413 TensorOperations,
415 NeuralOperations,
417 GroupedNeuralOperations,
419 HyperNeuralOperations,
421 PayloadMaterialization,
423 LogitsProcessing,
425 Sampling,
427 Randomness,
429 StateStorage,
431 StorageResidency,
433 ExactCompletion,
435 Observation,
437 Timing,
439 QueueBinding,
441 Communication,
443 Agreement,
445 Publication,
447 SameDeviceHandoff,
449 CrossDeviceTransfer,
451}
452
453impl SpeculativeMechanism {
454 const BASE: [Self; 13] = [
455 Self::TensorOperations,
456 Self::NeuralOperations,
457 Self::PayloadMaterialization,
458 Self::LogitsProcessing,
459 Self::Sampling,
460 Self::Randomness,
461 Self::StateStorage,
462 Self::StorageResidency,
463 Self::ExactCompletion,
464 Self::Observation,
465 Self::QueueBinding,
466 Self::Agreement,
467 Self::Publication,
468 ];
469}
470
471#[derive(Debug, Clone, Eq, PartialEq)]
473pub struct SpeculativeMechanismRequirements {
474 mechanisms: BTreeSet<SpeculativeMechanism>,
475}
476
477impl SpeculativeMechanismRequirements {
478 pub fn new(additional: impl IntoIterator<Item = SpeculativeMechanism>) -> Self {
480 Self {
481 mechanisms: SpeculativeMechanism::BASE
482 .into_iter()
483 .chain(additional)
484 .collect(),
485 }
486 }
487
488 pub const fn mechanisms(&self) -> &BTreeSet<SpeculativeMechanism> {
490 &self.mechanisms
491 }
492}
493
494impl Default for SpeculativeMechanismRequirements {
495 fn default() -> Self {
496 Self::new([])
497 }
498}
499
500#[derive(Debug, Clone, Eq, PartialEq)]
502pub struct SpeculativeMechanismCapabilities {
503 mechanisms: BTreeSet<SpeculativeMechanism>,
504}
505
506impl SpeculativeMechanismCapabilities {
507 pub fn new(mechanisms: impl IntoIterator<Item = SpeculativeMechanism>) -> Self {
509 Self {
510 mechanisms: mechanisms.into_iter().collect(),
511 }
512 }
513
514 pub const fn mechanisms(&self) -> &BTreeSet<SpeculativeMechanism> {
516 &self.mechanisms
517 }
518
519 pub fn supports(&self, mechanism: SpeculativeMechanism) -> bool {
521 self.mechanisms.contains(&mechanism)
522 }
523}
524
525#[derive(Debug, Clone, Copy, Eq, PartialEq)]
527#[non_exhaustive]
528pub enum SpeculativePlacementRequest {
529 Single,
531 SameDeviceSplit,
533 CrossDeviceSplit,
535}
536
537impl SpeculativePlacementRequest {
538 pub fn from_topology(
540 topology: SpeculativeExecutionTopology,
541 ) -> Result<Self, SpeculativeContractError> {
542 match topology {
543 SpeculativeExecutionTopology::Single => Ok(Self::Single),
544 SpeculativeExecutionTopology::SameDeviceSplit => Ok(Self::SameDeviceSplit),
545 SpeculativeExecutionTopology::CrossDeviceSplit => Ok(Self::CrossDeviceSplit),
546 _ => Err(SpeculativeContractError::new(
547 "unsupported speculative execution topology",
548 )),
549 }
550 }
551
552 fn topology(self) -> SpeculativeExecutionTopology {
553 match self {
554 Self::Single => SpeculativeExecutionTopology::Single,
555 Self::SameDeviceSplit => SpeculativeExecutionTopology::SameDeviceSplit,
556 Self::CrossDeviceSplit => SpeculativeExecutionTopology::CrossDeviceSplit,
557 }
558 }
559
560 fn required_mechanism(self) -> Option<SpeculativeMechanism> {
561 match self {
562 Self::Single => None,
563 Self::SameDeviceSplit => Some(SpeculativeMechanism::SameDeviceHandoff),
564 Self::CrossDeviceSplit => Some(SpeculativeMechanism::CrossDeviceTransfer),
565 }
566 }
567}
568
569#[derive(Debug, Clone, Copy, Eq, PartialEq)]
571pub struct SelectedSpeculativePlacement {
572 topology: SpeculativeExecutionTopology,
573}
574
575impl SelectedSpeculativePlacement {
576 pub const fn topology(self) -> SpeculativeExecutionTopology {
578 self.topology
579 }
580
581 pub const fn is_split(self) -> bool {
583 !matches!(self.topology, SpeculativeExecutionTopology::Single)
584 }
585
586 pub const fn crosses_devices(self) -> bool {
588 matches!(
589 self.topology,
590 SpeculativeExecutionTopology::CrossDeviceSplit
591 )
592 }
593}
594
595#[derive(Debug, Clone, Eq, PartialEq)]
597pub struct SpeculativeStateCacheIdentityIngredients {
598 target: SpeculativeIdentity,
599 strategy: SpeculativeIdentity,
600 assistant: Option<SpeculativeIdentity>,
601 tokenizer: Option<[u8; 32]>,
602 artifact: SpeculativeIdentity,
603 format: SpeculativeIdentity,
604 model_topology: SpeculativeIdentity,
605 rank: usize,
606 processor: SpeculativeIdentity,
607 state_components: Vec<SpeculativeIdentity>,
608}
609
610impl SpeculativeStateCacheIdentityIngredients {
611 #[allow(clippy::too_many_arguments)]
613 pub fn new(
614 target: SpeculativeIdentity,
615 strategy: SpeculativeIdentity,
616 assistant: Option<SpeculativeIdentity>,
617 tokenizer: Option<[u8; 32]>,
618 artifact: SpeculativeIdentity,
619 format: SpeculativeIdentity,
620 model_topology: SpeculativeIdentity,
621 rank: usize,
622 processor: SpeculativeIdentity,
623 state_components: Vec<SpeculativeIdentity>,
624 ) -> Result<Self, SpeculativeContractError> {
625 if state_components.is_empty() {
626 return Err(SpeculativeContractError::new(
627 "speculative state identity must contain at least one component",
628 ));
629 }
630 if state_components.iter().collect::<BTreeSet<_>>().len() != state_components.len() {
631 return Err(SpeculativeContractError::new(
632 "speculative state component identities must be unique",
633 ));
634 }
635 Ok(Self {
636 target,
637 strategy,
638 assistant,
639 tokenizer,
640 artifact,
641 format,
642 model_topology,
643 rank,
644 processor,
645 state_components,
646 })
647 }
648
649 pub const fn target(&self) -> &SpeculativeIdentity {
651 &self.target
652 }
653
654 pub const fn strategy(&self) -> &SpeculativeIdentity {
656 &self.strategy
657 }
658
659 pub const fn assistant(&self) -> Option<&SpeculativeIdentity> {
661 self.assistant.as_ref()
662 }
663
664 pub const fn tokenizer(&self) -> Option<[u8; 32]> {
666 self.tokenizer
667 }
668
669 pub const fn artifact(&self) -> &SpeculativeIdentity {
671 &self.artifact
672 }
673
674 pub const fn format(&self) -> &SpeculativeIdentity {
676 &self.format
677 }
678
679 pub const fn model_topology(&self) -> &SpeculativeIdentity {
681 &self.model_topology
682 }
683
684 pub const fn rank(&self) -> usize {
686 self.rank
687 }
688
689 pub const fn processor(&self) -> &SpeculativeIdentity {
691 &self.processor
692 }
693
694 pub fn state_components(&self) -> &[SpeculativeIdentity] {
696 &self.state_components
697 }
698}
699
700#[derive(Debug, Clone, Eq, PartialEq)]
702pub struct SpeculativeArchitectureCompatibilityProof {
703 target: SpeculativeIdentity,
704 strategy: SpeculativeIdentity,
705 capture_schema: SpeculativeIdentity,
706}
707
708impl SpeculativeArchitectureCompatibilityProof {
709 pub const fn new(
711 target: SpeculativeIdentity,
712 strategy: SpeculativeIdentity,
713 capture_schema: SpeculativeIdentity,
714 ) -> Self {
715 Self {
716 target,
717 strategy,
718 capture_schema,
719 }
720 }
721}
722
723#[derive(Debug, Clone, Eq, PartialEq)]
725pub struct SpeculativeRealizationRequirements {
726 target: SpeculativeIdentity,
727 strategy: SpeculativeStrategyRequirements,
728 capture: SpeculativeCaptureSchema,
729 mechanisms: SpeculativeMechanismRequirements,
730 state: SpeculativeStateCacheIdentityIngredients,
731}
732
733impl SpeculativeRealizationRequirements {
734 pub fn new(
736 target: SpeculativeIdentity,
737 strategy: SpeculativeStrategyRequirements,
738 capture: SpeculativeCaptureSchema,
739 mechanisms: SpeculativeMechanismRequirements,
740 state: SpeculativeStateCacheIdentityIngredients,
741 ) -> Result<Self, SpeculativeContractError> {
742 if state.target != target || state.strategy != strategy.identity {
743 return Err(SpeculativeContractError::new(
744 "speculative state identity differs from target or strategy requirements",
745 ));
746 }
747 if state.tokenizer != strategy.tokenizer_fingerprint {
748 return Err(SpeculativeContractError::new(
749 "speculative tokenizer identity differs between strategy and cache state",
750 ));
751 }
752 match strategy.class {
753 SpeculativeStrategyClass::External if state.assistant.is_none() => {
754 return Err(SpeculativeContractError::new(
755 "external strategy requires an assistant cache identity",
756 ));
757 }
758 SpeculativeStrategyClass::EmbeddedSequential
759 | SpeculativeStrategyClass::EmbeddedFused
760 if state.assistant.is_some() =>
761 {
762 return Err(SpeculativeContractError::new(
763 "embedded strategy cannot carry an external assistant identity",
764 ));
765 }
766 _ => {}
767 }
768 Ok(Self {
769 target,
770 strategy,
771 capture,
772 mechanisms,
773 state,
774 })
775 }
776
777 pub const fn target(&self) -> &SpeculativeIdentity {
779 &self.target
780 }
781
782 pub const fn strategy(&self) -> &SpeculativeStrategyRequirements {
784 &self.strategy
785 }
786
787 pub const fn capture(&self) -> &SpeculativeCaptureSchema {
789 &self.capture
790 }
791
792 pub const fn mechanisms(&self) -> &SpeculativeMechanismRequirements {
794 &self.mechanisms
795 }
796
797 pub const fn state(&self) -> &SpeculativeStateCacheIdentityIngredients {
799 &self.state
800 }
801}
802
803#[derive(Debug, Clone, Eq, PartialEq)]
805pub struct SpeculativeSelectionRequest {
806 placement: SpeculativePlacementRequest,
807 architecture: Option<SpeculativeArchitectureCompatibilityProof>,
808 tokenizer: Option<TokenizerCompatibilityProof>,
809 capture: SpeculativeCaptureSchema,
810}
811
812impl SpeculativeSelectionRequest {
813 pub const fn new(
815 placement: SpeculativePlacementRequest,
816 capture: SpeculativeCaptureSchema,
817 ) -> Self {
818 Self {
819 placement,
820 architecture: None,
821 tokenizer: None,
822 capture,
823 }
824 }
825
826 pub fn with_architecture_proof(
828 mut self,
829 proof: SpeculativeArchitectureCompatibilityProof,
830 ) -> Self {
831 self.architecture = Some(proof);
832 self
833 }
834
835 pub fn with_tokenizer_proof(mut self, proof: TokenizerCompatibilityProof) -> Self {
837 self.tokenizer = Some(proof);
838 self
839 }
840}
841
842#[derive(Debug, Clone, Copy, Eq, PartialEq)]
844pub struct SelectedSpeculativeSampling;
845
846#[derive(Debug, Clone, Copy, Eq, PartialEq)]
848pub struct SelectedSpeculativeCompletion;
849
850#[derive(Debug, Clone, Eq, PartialEq)]
852pub struct SelectedSpeculativeState {
853 identity: SpeculativeStateCacheIdentityIngredients,
854 placement: SpeculativeExecutionTopology,
855}
856
857#[derive(Debug, Clone, Eq, PartialEq)]
864pub struct SpeculativeLaneIdentity {
865 realization: SelectedSpeculativeState,
866 prepared_input: SpeculativeIdentity,
867 capture_generation: u64,
868}
869
870impl SpeculativeLaneIdentity {
871 pub const fn realization(&self) -> &SelectedSpeculativeState {
873 &self.realization
874 }
875
876 pub const fn prepared_input(&self) -> &SpeculativeIdentity {
878 &self.prepared_input
879 }
880
881 pub const fn capture_generation(&self) -> u64 {
883 self.capture_generation
884 }
885}
886
887impl SelectedSpeculativeState {
888 pub const fn identity(&self) -> &SpeculativeStateCacheIdentityIngredients {
890 &self.identity
891 }
892
893 pub const fn placement(&self) -> SpeculativeExecutionTopology {
895 self.placement
896 }
897}
898
899#[derive(Debug, Clone, Eq, PartialEq)]
901pub struct SelectedSpeculativeRealization {
902 requirements: SpeculativeRealizationRequirements,
903 placement: SelectedSpeculativePlacement,
904 sampling: SelectedSpeculativeSampling,
905 state: SelectedSpeculativeState,
906 completion: SelectedSpeculativeCompletion,
907}
908
909impl SelectedSpeculativeRealization {
910 pub const fn requirements(&self) -> &SpeculativeRealizationRequirements {
912 &self.requirements
913 }
914
915 pub const fn placement(&self) -> SelectedSpeculativePlacement {
917 self.placement
918 }
919
920 pub const fn sampling(&self) -> SelectedSpeculativeSampling {
922 self.sampling
923 }
924
925 pub const fn state(&self) -> &SelectedSpeculativeState {
927 &self.state
928 }
929
930 pub const fn completion(&self) -> SelectedSpeculativeCompletion {
932 self.completion
933 }
934
935 pub fn lane_identity(
937 &self,
938 prepared_input: SpeculativeIdentity,
939 capture_generation: u64,
940 ) -> SpeculativeLaneIdentity {
941 SpeculativeLaneIdentity {
942 realization: self.state.clone(),
943 prepared_input,
944 capture_generation,
945 }
946 }
947
948 pub fn validate_capture<T>(
950 &self,
951 lane: &SpeculativeLaneIdentity,
952 capture: &SpeculativeCaptureEnvelope<T>,
953 ) -> Result<(), SpeculativeCaptureError> {
954 if lane.realization != self.state {
955 return Err(SpeculativeCaptureError::RealizationMismatch);
956 }
957 capture.validate_against(&self.requirements.capture, lane.capture_generation)
958 }
959}
960
961#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
963#[error("speculative realization is unsupported: {issues}", issues = .issues.join("; "))]
964pub struct SpeculativeSelectionError {
965 issues: Vec<String>,
966}
967
968impl SpeculativeSelectionError {
969 pub fn issues(&self) -> &[String] {
971 &self.issues
972 }
973}
974
975pub fn select_speculative_realization(
977 requirements: &SpeculativeRealizationRequirements,
978 request: &SpeculativeSelectionRequest,
979 capabilities: &SpeculativeMechanismCapabilities,
980) -> Result<SelectedSpeculativeRealization, SpeculativeSelectionError> {
981 let mut issues = Vec::new();
982 match &request.architecture {
983 Some(proof)
984 if proof.target == requirements.target
985 && proof.strategy == requirements.strategy.identity
986 && proof.capture_schema == requirements.capture.identity => {}
987 Some(_) => issues.push("architecture compatibility proof identity mismatch".into()),
988 None => issues.push("architecture compatibility proof is missing".into()),
989 }
990 match (
991 requirements.strategy.tokenizer_fingerprint,
992 request.tokenizer.as_ref(),
993 ) {
994 (Some(expected), Some(proof)) if proof.fingerprint() == expected => {}
995 (Some(_), Some(_)) => issues.push("tokenizer compatibility proof identity mismatch".into()),
996 (Some(_), None) => issues.push("tokenizer compatibility proof is missing".into()),
997 (None, Some(_)) => {
998 issues.push("embedded strategy received an external tokenizer proof".into())
999 }
1000 (None, None) => {}
1001 }
1002 if request.capture != requirements.capture {
1003 issues.push("target capture path, order, shape, owner, or observation mismatch".into());
1004 }
1005 for mechanism in requirements.mechanisms.mechanisms() {
1006 if !capabilities.supports(*mechanism) {
1007 issues.push(format!("missing speculative mechanism {mechanism:?}"));
1008 }
1009 }
1010 if let Some(mechanism) = request.placement.required_mechanism() {
1011 if !capabilities.supports(mechanism) {
1012 issues.push(format!(
1013 "placement {:?} requires speculative mechanism {mechanism:?}",
1014 request.placement
1015 ));
1016 }
1017 }
1018 if !issues.is_empty() {
1019 return Err(SpeculativeSelectionError { issues });
1020 }
1021 let placement = SelectedSpeculativePlacement {
1022 topology: request.placement.topology(),
1023 };
1024 Ok(SelectedSpeculativeRealization {
1025 requirements: requirements.clone(),
1026 placement,
1027 sampling: SelectedSpeculativeSampling,
1028 state: SelectedSpeculativeState {
1029 identity: requirements.state.clone(),
1030 placement: placement.topology,
1031 },
1032 completion: SelectedSpeculativeCompletion,
1033 })
1034}
1035
1036#[derive(Debug)]
1038pub struct ConstructedSpeculativeResources<P, M, S, Q, T> {
1039 payload: P,
1040 modules: M,
1041 state: S,
1042 queues: Q,
1043 transfer: Option<T>,
1044}
1045
1046impl<P, M, S, Q, T> ConstructedSpeculativeResources<P, M, S, Q, T> {
1047 pub fn into_parts(self) -> (P, M, S, Q, Option<T>) {
1049 (
1050 self.payload,
1051 self.modules,
1052 self.state,
1053 self.queues,
1054 self.transfer,
1055 )
1056 }
1057}
1058
1059#[derive(Debug)]
1061pub struct PreparedSpeculativeRealization<P, M, S, Q, T> {
1062 selected: SelectedSpeculativeRealization,
1063 resources: ConstructedSpeculativeResources<P, M, S, Q, T>,
1064}
1065
1066impl<P, M, S, Q, T> PreparedSpeculativeRealization<P, M, S, Q, T> {
1067 pub const fn selected(&self) -> &SelectedSpeculativeRealization {
1069 &self.selected
1070 }
1071
1072 pub fn into_parts(
1074 self,
1075 ) -> (
1076 SelectedSpeculativeRealization,
1077 ConstructedSpeculativeResources<P, M, S, Q, T>,
1078 ) {
1079 (self.selected, self.resources)
1080 }
1081}
1082
1083#[derive(Debug, thiserror::Error)]
1085pub enum SpeculativePreparationError<E> {
1086 #[error(transparent)]
1088 Lifecycle(#[from] SpeculativeOutputError),
1089 #[error(transparent)]
1091 Selection(#[from] SpeculativeSelectionError),
1092 #[error("speculative payload construction failed")]
1094 Payload(#[source] E),
1095 #[error("speculative module construction failed")]
1097 Modules(#[source] E),
1098 #[error("speculative state construction failed")]
1100 State(#[source] E),
1101 #[error("speculative queue construction failed")]
1103 Queues(#[source] E),
1104 #[error("speculative transfer construction failed")]
1106 Transfer(#[source] E),
1107}
1108
1109#[allow(clippy::too_many_arguments)]
1111pub fn select_and_prepare_speculative_realization<P, M, S, Q, T, E>(
1112 requirements: &SpeculativeRealizationRequirements,
1113 request: &SpeculativeSelectionRequest,
1114 capabilities: &SpeculativeMechanismCapabilities,
1115 payload: impl FnOnce(&SelectedSpeculativeRealization) -> Result<P, E>,
1116 modules: impl FnOnce(&SelectedSpeculativeRealization, &P) -> Result<M, E>,
1117 state: impl FnOnce(&SelectedSpeculativeRealization, &M) -> Result<S, E>,
1118 queues: impl FnOnce(&SelectedSpeculativeRealization) -> Result<Q, E>,
1119 transfer: impl FnOnce(&SelectedSpeculativeRealization, &Q) -> Result<T, E>,
1120) -> Result<PreparedSpeculativeRealization<P, M, S, Q, T>, SpeculativePreparationError<E>> {
1121 select_and_prepare_speculative_realization_observed(
1122 requirements,
1123 request,
1124 capabilities,
1125 &|_| Ok(()),
1126 payload,
1127 modules,
1128 state,
1129 queues,
1130 transfer,
1131 )
1132}
1133
1134#[allow(clippy::too_many_arguments)]
1140pub fn select_and_prepare_speculative_realization_observed<P, M, S, Q, T, E>(
1141 requirements: &SpeculativeRealizationRequirements,
1142 request: &SpeculativeSelectionRequest,
1143 capabilities: &SpeculativeMechanismCapabilities,
1144 observer: &impl SpeculativeLifecycleObserver,
1145 payload: impl FnOnce(&SelectedSpeculativeRealization) -> Result<P, E>,
1146 modules: impl FnOnce(&SelectedSpeculativeRealization, &P) -> Result<M, E>,
1147 state: impl FnOnce(&SelectedSpeculativeRealization, &M) -> Result<S, E>,
1148 queues: impl FnOnce(&SelectedSpeculativeRealization) -> Result<Q, E>,
1149 transfer: impl FnOnce(&SelectedSpeculativeRealization, &Q) -> Result<T, E>,
1150) -> Result<PreparedSpeculativeRealization<P, M, S, Q, T>, SpeculativePreparationError<E>> {
1151 observer.observe(SpeculativeLifecycleStage::Admission)?;
1152 let selected = select_speculative_realization(requirements, request, capabilities)?;
1153 observer.observe(SpeculativeLifecycleStage::Compatibility)?;
1154 observer.observe(SpeculativeLifecycleStage::Input)?;
1155 let payload = payload(&selected).map_err(SpeculativePreparationError::Payload)?;
1156 observer.observe(SpeculativeLifecycleStage::Execution)?;
1157 let modules = modules(&selected, &payload).map_err(SpeculativePreparationError::Modules)?;
1158 let state = state(&selected, &modules).map_err(SpeculativePreparationError::State)?;
1159 let queues = queues(&selected).map_err(SpeculativePreparationError::Queues)?;
1160 let transfer = if selected.placement.crosses_devices() {
1161 observer.observe(SpeculativeLifecycleStage::Transfer)?;
1162 Some(transfer(&selected, &queues).map_err(SpeculativePreparationError::Transfer)?)
1163 } else {
1164 None
1165 };
1166 Ok(PreparedSpeculativeRealization {
1167 selected,
1168 resources: ConstructedSpeculativeResources {
1169 payload,
1170 modules,
1171 state,
1172 queues,
1173 transfer,
1174 },
1175 })
1176}
1177
1178#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1180#[error("invalid speculative contract: {message}")]
1181pub struct SpeculativeContractError {
1182 message: String,
1183}
1184
1185impl SpeculativeContractError {
1186 fn new(message: impl Into<String>) -> Self {
1187 Self {
1188 message: message.into(),
1189 }
1190 }
1191}
1192
1193#[cfg(test)]
1194mod tests {
1195 use std::{
1196 cell::{Cell, RefCell},
1197 convert::Infallible,
1198 rc::Rc,
1199 sync::{
1200 atomic::{AtomicUsize, Ordering},
1201 Arc, Mutex,
1202 },
1203 };
1204
1205 use super::*;
1206
1207 fn id(value: &str) -> SpeculativeIdentity {
1208 SpeculativeIdentity::new(value).unwrap()
1209 }
1210
1211 fn capture() -> SpeculativeCaptureSchema {
1212 SpeculativeCaptureSchema::new(
1213 id("capture-v1"),
1214 [
1215 SpeculativeCaptureEntry::new(
1216 id("layers.1.output"),
1217 vec![1, 2, 8],
1218 id("output-rank"),
1219 id("layer-1-seam"),
1220 )
1221 .unwrap(),
1222 SpeculativeCaptureEntry::new(
1223 id("layers.3.output"),
1224 vec![1, 2, 8],
1225 id("output-rank"),
1226 id("layer-3-seam"),
1227 )
1228 .unwrap(),
1229 ],
1230 )
1231 .unwrap()
1232 }
1233
1234 fn requirements(class: SpeculativeStrategyClass) -> SpeculativeRealizationRequirements {
1235 requirements_with_mechanisms(class, [])
1236 }
1237
1238 fn requirements_with_mechanisms(
1239 class: SpeculativeStrategyClass,
1240 additional_mechanisms: impl IntoIterator<Item = SpeculativeMechanism>,
1241 ) -> SpeculativeRealizationRequirements {
1242 let target = id("target-v1");
1243 let strategy_id = id("strategy-v1");
1244 let (strategy, assistant, tokenizer) = match class {
1245 SpeculativeStrategyClass::External => (
1246 SpeculativeStrategyRequirements::external(
1247 strategy_id.clone(),
1248 NonZeroUsize::new(4).unwrap(),
1249 [7; 32],
1250 ),
1251 Some(id("assistant-v1")),
1252 Some([7; 32]),
1253 ),
1254 _ => (
1255 SpeculativeStrategyRequirements::embedded(
1256 class,
1257 strategy_id.clone(),
1258 NonZeroUsize::new(4).unwrap(),
1259 )
1260 .unwrap(),
1261 None,
1262 None,
1263 ),
1264 };
1265 let state = SpeculativeStateCacheIdentityIngredients::new(
1266 target.clone(),
1267 strategy_id,
1268 assistant,
1269 tokenizer,
1270 id("artifact-v1"),
1271 id("safetensors"),
1272 id("tp2-rank0"),
1273 0,
1274 id("text-processor-v1"),
1275 vec![id("target-state"), id("prediction-state")],
1276 )
1277 .unwrap();
1278 SpeculativeRealizationRequirements::new(
1279 target,
1280 strategy,
1281 capture(),
1282 SpeculativeMechanismRequirements::new(additional_mechanisms),
1283 state,
1284 )
1285 .unwrap()
1286 }
1287
1288 fn capabilities() -> SpeculativeMechanismCapabilities {
1289 SpeculativeMechanismCapabilities::new(SpeculativeMechanism::BASE.into_iter().chain([
1290 SpeculativeMechanism::GroupedNeuralOperations,
1291 SpeculativeMechanism::HyperNeuralOperations,
1292 SpeculativeMechanism::Timing,
1293 SpeculativeMechanism::Communication,
1294 SpeculativeMechanism::SameDeviceHandoff,
1295 SpeculativeMechanism::CrossDeviceTransfer,
1296 ]))
1297 }
1298
1299 fn request(
1300 requirements: &SpeculativeRealizationRequirements,
1301 placement: SpeculativePlacementRequest,
1302 ) -> SpeculativeSelectionRequest {
1303 let mut request = SpeculativeSelectionRequest::new(placement, requirements.capture.clone())
1304 .with_architecture_proof(SpeculativeArchitectureCompatibilityProof::new(
1305 requirements.target.clone(),
1306 requirements.strategy.identity.clone(),
1307 requirements.capture.identity.clone(),
1308 ));
1309 if let Some(tokenizer) = requirements.strategy.tokenizer_fingerprint {
1310 request = request.with_tokenizer_proof(
1311 TokenizerCompatibilityProof::prove(tokenizer, tokenizer).unwrap(),
1312 );
1313 }
1314 request
1315 }
1316
1317 fn construction_counter() -> (Rc<Cell<usize>>, impl Fn() -> usize) {
1318 let counter = Rc::new(Cell::new(0));
1319 let read = Rc::clone(&counter);
1320 (counter, move || read.get())
1321 }
1322
1323 #[test]
1324 fn exact_selection_precedes_every_construction_callback() {
1325 let requirements = requirements(SpeculativeStrategyClass::External);
1326 let request = request(&requirements, SpeculativePlacementRequest::CrossDeviceSplit);
1327 let (counter, calls) = construction_counter();
1328 let prepared = select_and_prepare_speculative_realization(
1329 &requirements,
1330 &request,
1331 &capabilities(),
1332 {
1333 let counter = Rc::clone(&counter);
1334 move |_| {
1335 assert_eq!(counter.get(), 0);
1336 counter.set(1);
1337 Ok::<_, Infallible>("payload")
1338 }
1339 },
1340 {
1341 let counter = Rc::clone(&counter);
1342 move |_, payload| {
1343 assert_eq!((*payload, counter.get()), ("payload", 1));
1344 counter.set(2);
1345 Ok::<_, Infallible>("modules")
1346 }
1347 },
1348 {
1349 let counter = Rc::clone(&counter);
1350 move |_, modules| {
1351 assert_eq!((*modules, counter.get()), ("modules", 2));
1352 counter.set(3);
1353 Ok::<_, Infallible>("state")
1354 }
1355 },
1356 {
1357 let counter = Rc::clone(&counter);
1358 move |_| {
1359 assert_eq!(counter.get(), 3);
1360 counter.set(4);
1361 Ok::<_, Infallible>("queues")
1362 }
1363 },
1364 {
1365 let counter = Rc::clone(&counter);
1366 move |selected, queues| {
1367 assert!(selected.placement().crosses_devices());
1368 assert_eq!((*queues, counter.get()), ("queues", 4));
1369 counter.set(5);
1370 Ok::<_, Infallible>("transfer")
1371 }
1372 },
1373 )
1374 .unwrap();
1375 assert_eq!(calls(), 5);
1376 assert_eq!(
1377 prepared.selected().state().placement(),
1378 SpeculativeExecutionTopology::CrossDeviceSplit
1379 );
1380 let (_, resources) = prepared.into_parts();
1381 assert_eq!(
1382 resources.into_parts(),
1383 ("payload", "modules", "state", "queues", Some("transfer"))
1384 );
1385 }
1386
1387 #[test]
1388 fn production_lifecycle_observation_precedes_each_construction_boundary() {
1389 let requirements = requirements(SpeculativeStrategyClass::External);
1390 let request = request(&requirements, SpeculativePlacementRequest::CrossDeviceSplit);
1391 let observed = Arc::new(Mutex::new(Vec::new()));
1392 let observer = {
1393 let observed = Arc::clone(&observed);
1394 move |stage| {
1395 observed.lock().unwrap().push(stage);
1396 Ok(())
1397 }
1398 };
1399 select_and_prepare_speculative_realization_observed(
1400 &requirements,
1401 &request,
1402 &capabilities(),
1403 &observer,
1404 |_| Ok::<_, Infallible>("payload"),
1405 |_, _| Ok("modules"),
1406 |_, _| Ok("state"),
1407 |_| Ok("queues"),
1408 |_, _| Ok("transfer"),
1409 )
1410 .unwrap();
1411 assert_eq!(
1412 *observed.lock().unwrap(),
1413 [
1414 SpeculativeLifecycleStage::Admission,
1415 SpeculativeLifecycleStage::Compatibility,
1416 SpeculativeLifecycleStage::Input,
1417 SpeculativeLifecycleStage::Execution,
1418 SpeculativeLifecycleStage::Transfer,
1419 ]
1420 );
1421 }
1422
1423 #[test]
1424 fn lifecycle_failure_prevents_its_native_construction_boundary() {
1425 let requirements = requirements(SpeculativeStrategyClass::External);
1426 let request = request(&requirements, SpeculativePlacementRequest::CrossDeviceSplit);
1427 for (failure, expected_callbacks) in [
1428 (SpeculativeLifecycleStage::Admission, 0),
1429 (SpeculativeLifecycleStage::Compatibility, 0),
1430 (SpeculativeLifecycleStage::Input, 0),
1431 (SpeculativeLifecycleStage::Execution, 1),
1432 (SpeculativeLifecycleStage::Transfer, 4),
1433 ] {
1434 let callbacks = Arc::new(AtomicUsize::new(0));
1435 let observer = move |stage| {
1436 if stage == failure {
1437 Err(SpeculativeOutputError::semantic(
1438 "lifecycle observation",
1439 format!("failed at {stage:?}"),
1440 ))
1441 } else {
1442 Ok(())
1443 }
1444 };
1445 let increment = || {
1446 callbacks.fetch_add(1, Ordering::SeqCst);
1447 Ok::<_, Infallible>(())
1448 };
1449 let result = select_and_prepare_speculative_realization_observed(
1450 &requirements,
1451 &request,
1452 &capabilities(),
1453 &observer,
1454 |_| increment(),
1455 |_, _| increment(),
1456 |_, _| increment(),
1457 |_| increment(),
1458 |_, _| increment(),
1459 );
1460 assert!(matches!(
1461 result,
1462 Err(SpeculativePreparationError::Lifecycle(_))
1463 ));
1464 assert_eq!(callbacks.load(Ordering::SeqCst), expected_callbacks);
1465 }
1466
1467 let observed = Arc::new(Mutex::new(Vec::new()));
1468 let observer = {
1469 let observed = Arc::clone(&observed);
1470 move |stage| {
1471 observed.lock().unwrap().push(stage);
1472 Ok(())
1473 }
1474 };
1475 let mut incompatible = request.clone();
1476 incompatible.architecture = None;
1477 let result = select_and_prepare_speculative_realization_observed(
1478 &requirements,
1479 &incompatible,
1480 &capabilities(),
1481 &observer,
1482 |_| Ok::<_, Infallible>(()),
1483 |_, _| Ok(()),
1484 |_, _| Ok(()),
1485 |_| Ok(()),
1486 |_, _| Ok(()),
1487 );
1488 assert!(matches!(
1489 result,
1490 Err(SpeculativePreparationError::Selection(_))
1491 ));
1492 assert_eq!(
1493 *observed.lock().unwrap(),
1494 [SpeculativeLifecycleStage::Admission]
1495 );
1496 }
1497
1498 #[test]
1499 fn missing_mechanisms_and_placement_fail_before_native_work() {
1500 let requirements = requirements_with_mechanisms(
1501 SpeculativeStrategyClass::EmbeddedSequential,
1502 [
1503 SpeculativeMechanism::GroupedNeuralOperations,
1504 SpeculativeMechanism::HyperNeuralOperations,
1505 SpeculativeMechanism::Timing,
1506 SpeculativeMechanism::Communication,
1507 ],
1508 );
1509 let cases = SpeculativeMechanism::BASE
1510 .into_iter()
1511 .chain([
1512 SpeculativeMechanism::GroupedNeuralOperations,
1513 SpeculativeMechanism::HyperNeuralOperations,
1514 SpeculativeMechanism::Timing,
1515 SpeculativeMechanism::Communication,
1516 ])
1517 .map(|mechanism| (mechanism, SpeculativePlacementRequest::Single))
1518 .chain([
1519 (
1520 SpeculativeMechanism::SameDeviceHandoff,
1521 SpeculativePlacementRequest::SameDeviceSplit,
1522 ),
1523 (
1524 SpeculativeMechanism::CrossDeviceTransfer,
1525 SpeculativePlacementRequest::CrossDeviceSplit,
1526 ),
1527 ]);
1528 for (missing, placement) in cases {
1529 let supported = capabilities();
1530 let available = supported
1531 .mechanisms()
1532 .iter()
1533 .copied()
1534 .filter(|mechanism| *mechanism != missing);
1535 let (counter, calls) = construction_counter();
1536 let result = select_and_prepare_speculative_realization(
1537 &requirements,
1538 &request(&requirements, placement),
1539 &SpeculativeMechanismCapabilities::new(available),
1540 {
1541 let counter = Rc::clone(&counter);
1542 move |_| {
1543 counter.set(counter.get() + 1);
1544 Ok::<_, Infallible>(())
1545 }
1546 },
1547 |_, _| Ok(()),
1548 |_, _| Ok(()),
1549 |_| Ok(()),
1550 |_, _| Ok(()),
1551 );
1552 assert!(matches!(
1553 result,
1554 Err(SpeculativePreparationError::Selection(_))
1555 ));
1556 assert_eq!(calls(), 0, "missing {missing:?} reached payload work");
1557 }
1558 }
1559
1560 #[test]
1561 fn construction_failure_stops_each_later_stage() {
1562 let requirements = requirements(SpeculativeStrategyClass::External);
1563 let request = request(&requirements, SpeculativePlacementRequest::CrossDeviceSplit);
1564 for (failure, expected_work) in [
1565 ("payload", &["payload"][..]),
1566 ("modules", &["payload", "modules"][..]),
1567 ("state", &["payload", "modules", "state"][..]),
1568 ("queues", &["payload", "modules", "state", "queues"][..]),
1569 (
1570 "transfer",
1571 &["payload", "modules", "state", "queues", "transfer"][..],
1572 ),
1573 ] {
1574 let trace = Rc::new(RefCell::new(Vec::new()));
1575 let result = select_and_prepare_speculative_realization(
1576 &requirements,
1577 &request,
1578 &capabilities(),
1579 {
1580 let trace = Rc::clone(&trace);
1581 move |_| {
1582 trace.borrow_mut().push("payload");
1583 if failure == "payload" {
1584 Err("payload failure")
1585 } else {
1586 Ok("payload")
1587 }
1588 }
1589 },
1590 {
1591 let trace = Rc::clone(&trace);
1592 move |_, _| {
1593 trace.borrow_mut().push("modules");
1594 if failure == "modules" {
1595 Err("modules failure")
1596 } else {
1597 Ok("modules")
1598 }
1599 }
1600 },
1601 {
1602 let trace = Rc::clone(&trace);
1603 move |_, _| {
1604 trace.borrow_mut().push("state");
1605 if failure == "state" {
1606 Err("state failure")
1607 } else {
1608 Ok("state")
1609 }
1610 }
1611 },
1612 {
1613 let trace = Rc::clone(&trace);
1614 move |_| {
1615 trace.borrow_mut().push("queues");
1616 if failure == "queues" {
1617 Err("queues failure")
1618 } else {
1619 Ok("queues")
1620 }
1621 }
1622 },
1623 {
1624 let trace = Rc::clone(&trace);
1625 move |_, _| {
1626 trace.borrow_mut().push("transfer");
1627 if failure == "transfer" {
1628 Err("transfer failure")
1629 } else {
1630 Ok("transfer")
1631 }
1632 }
1633 },
1634 );
1635 assert_eq!(&*trace.borrow(), expected_work, "failure at {failure}");
1636 assert!(
1637 matches!(
1638 (failure, result),
1639 (
1640 "payload",
1641 Err(SpeculativePreparationError::Payload("payload failure"))
1642 ) | (
1643 "modules",
1644 Err(SpeculativePreparationError::Modules("modules failure"))
1645 ) | (
1646 "state",
1647 Err(SpeculativePreparationError::State("state failure"))
1648 ) | (
1649 "queues",
1650 Err(SpeculativePreparationError::Queues("queues failure"))
1651 ) | (
1652 "transfer",
1653 Err(SpeculativePreparationError::Transfer("transfer failure"))
1654 )
1655 ),
1656 "unexpected result for {failure}"
1657 );
1658 }
1659 }
1660
1661 #[test]
1662 fn proof_and_capture_perturbations_fail_before_construction() {
1663 let requirements = requirements(SpeculativeStrategyClass::External);
1664 let base = request(&requirements, SpeculativePlacementRequest::Single);
1665 let mut reordered = requirements.capture.clone();
1666 reordered.entries.swap(0, 1);
1667 let mut wrong_shape = requirements.capture.clone();
1668 wrong_shape.entries[0].shape[2] += 1;
1669 let mut wrong_owner = requirements.capture.clone();
1670 wrong_owner.entries[0].owner = id("other-output-rank");
1671 let mut wrong_observation = requirements.capture.clone();
1672 wrong_observation.entries[0].observation = id("other-observation-seam");
1673 let requests = [
1674 SpeculativeSelectionRequest {
1675 architecture: None,
1676 ..base.clone()
1677 },
1678 SpeculativeSelectionRequest {
1679 architecture: Some(SpeculativeArchitectureCompatibilityProof::new(
1680 id("other-target"),
1681 requirements.strategy.identity.clone(),
1682 requirements.capture.identity.clone(),
1683 )),
1684 ..base.clone()
1685 },
1686 SpeculativeSelectionRequest {
1687 tokenizer: None,
1688 ..base.clone()
1689 },
1690 SpeculativeSelectionRequest {
1691 tokenizer: Some(TokenizerCompatibilityProof::prove([8; 32], [8; 32]).unwrap()),
1692 ..base.clone()
1693 },
1694 SpeculativeSelectionRequest {
1695 capture: reordered,
1696 ..base.clone()
1697 },
1698 SpeculativeSelectionRequest {
1699 capture: wrong_shape,
1700 ..base.clone()
1701 },
1702 SpeculativeSelectionRequest {
1703 capture: wrong_owner,
1704 ..base.clone()
1705 },
1706 SpeculativeSelectionRequest {
1707 capture: wrong_observation,
1708 ..base
1709 },
1710 ];
1711 for request in requests {
1712 let trace = Rc::new(RefCell::new(Vec::new()));
1713 let result = select_and_prepare_speculative_realization(
1714 &requirements,
1715 &request,
1716 &capabilities(),
1717 {
1718 let trace = Rc::clone(&trace);
1719 move |_| {
1720 trace.borrow_mut().push("payload");
1721 Ok::<_, Infallible>(())
1722 }
1723 },
1724 {
1725 let trace = Rc::clone(&trace);
1726 move |_, _| {
1727 trace.borrow_mut().push("modules");
1728 Ok(())
1729 }
1730 },
1731 {
1732 let trace = Rc::clone(&trace);
1733 move |_, _| {
1734 trace.borrow_mut().push("state");
1735 Ok(())
1736 }
1737 },
1738 {
1739 let trace = Rc::clone(&trace);
1740 move |_| {
1741 trace.borrow_mut().push("queues");
1742 Ok(())
1743 }
1744 },
1745 {
1746 let trace = Rc::clone(&trace);
1747 move |_, _| {
1748 trace.borrow_mut().push("transfer");
1749 Ok(())
1750 }
1751 },
1752 );
1753 assert!(matches!(
1754 result,
1755 Err(SpeculativePreparationError::Selection(_))
1756 ));
1757 assert!(trace.borrow().is_empty());
1758 }
1759 }
1760
1761 #[test]
1762 fn capture_envelope_validates_cardinality_schema_and_generation() {
1763 let schema = capture();
1764 let envelope = SpeculativeCaptureEnvelope::new(
1765 SpeculativeCaptureMetadata::new(schema.clone(), 11),
1766 vec![10, 30],
1767 )
1768 .unwrap();
1769 envelope.validate_against(&schema, 11).unwrap();
1770 assert!(matches!(
1771 envelope.validate_against(&schema, 12),
1772 Err(SpeculativeCaptureError::GenerationMismatch { .. })
1773 ));
1774 let mut reordered = schema.clone();
1775 reordered.entries.swap(0, 1);
1776 assert_eq!(
1777 envelope.validate_against(&reordered, 11),
1778 Err(SpeculativeCaptureError::SchemaMismatch)
1779 );
1780 assert!(matches!(
1781 SpeculativeCaptureEnvelope::new(SpeculativeCaptureMetadata::new(schema, 11), vec![10]),
1782 Err(SpeculativeCaptureError::ValueCount { .. })
1783 ));
1784 }
1785
1786 #[test]
1787 fn request_sized_capture_dimensions_close_exactly_at_the_lane_boundary() {
1788 let schema = SpeculativeCaptureSchema::new(
1789 id("bounded-capture-v1"),
1790 [SpeculativeCaptureEntry::new(
1791 id("layers.1.output"),
1792 vec![1, 128, 8],
1793 id("output-rank"),
1794 id("layer-1-seam"),
1795 )
1796 .unwrap()
1797 .with_bounded_dimension(1)
1798 .unwrap()],
1799 )
1800 .unwrap();
1801 let actual = schema.instantiate([vec![1, 7, 8]]).unwrap();
1802 let envelope =
1803 SpeculativeCaptureEnvelope::new(SpeculativeCaptureMetadata::new(actual, 4), vec![10])
1804 .unwrap();
1805 envelope.validate_against(&schema, 4).unwrap();
1806
1807 assert_eq!(
1808 schema.instantiate([vec![1, 129, 8]]),
1809 Err(SpeculativeCaptureError::ShapeMismatch)
1810 );
1811 assert_eq!(
1812 schema.instantiate([vec![1, 7, 9]]),
1813 Err(SpeculativeCaptureError::ShapeMismatch)
1814 );
1815 }
1816
1817 #[test]
1818 fn request_local_identity_is_closed_only_after_model_selection() {
1819 let requirements = requirements(SpeculativeStrategyClass::EmbeddedSequential);
1820 let selected = select_speculative_realization(
1821 &requirements,
1822 &request(&requirements, SpeculativePlacementRequest::Single),
1823 &capabilities(),
1824 )
1825 .unwrap();
1826 let lane = selected.lane_identity(id("prepared-input-42"), 11);
1827 let envelope = SpeculativeCaptureEnvelope::new(
1828 SpeculativeCaptureMetadata::new(requirements.capture.clone(), 11),
1829 vec![10, 30],
1830 )
1831 .unwrap();
1832 selected.validate_capture(&lane, &envelope).unwrap();
1833 assert_eq!(lane.prepared_input().as_str(), "prepared-input-42");
1834 assert_eq!(lane.capture_generation(), 11);
1835
1836 let other = select_speculative_realization(
1837 &requirements,
1838 &request(&requirements, SpeculativePlacementRequest::SameDeviceSplit),
1839 &capabilities(),
1840 )
1841 .unwrap();
1842 assert_eq!(
1843 other.validate_capture(&lane, &envelope),
1844 Err(SpeculativeCaptureError::RealizationMismatch)
1845 );
1846 }
1847}