Skip to main content

eredu_runtime/
replicated_text.rs

1//! Selection contracts for replicated text architectures.
2
3use std::{
4    collections::{BTreeMap, BTreeSet},
5    path::{Path, PathBuf},
6};
7
8use eredu_checkpoint::{LinearFormat, SourceTensorEncoding, StoredDtype};
9use eredu_core::{
10    cache::{StateComponentPolicy, StateTensorDtype},
11    checkpoint::TensorDtype,
12    ParallelTopology, QuantizationRequest, SessionCapabilities,
13};
14use eredu_nn::{NeuralBackend, NeuralOperatorCapabilities};
15
16use crate::{
17    ArchitectureGroupTransport, ArchitectureParameterDescription, ArchitecturePartition,
18    CacheResidencyPolicy, ExecutionGraph, ExecutionGroupId, ExecutionUnitLayout,
19    LayerWeightResidency, LayeredArchitecture, ParameterGroupOwner, ParameterGroupSpec,
20    RuntimeState, StateLayout,
21};
22
23/// Statically dispatched text-input seam for a layered decoder.
24///
25/// Routed, composite, partitioned, prediction, and realtime execution use
26/// separate extension contracts rather than adding requirements here.
27pub trait ReplicatedTextArchitecture<B, S>: LayeredArchitecture<B, S>
28where
29    B: NeuralBackend,
30    S: RuntimeState<B>,
31{
32    /// Forms the architecture-owned borrowed input for one text pass.
33    fn text_input<'a>(tokens: &'a B::Tensor, mask: Option<&'a B::Tensor>) -> Self::Input<'a>;
34
35    /// Declares how a causal-text session projects a complete architecture output.
36    fn text_output_selection(&self) -> ReplicatedTextOutputSelection {
37        ReplicatedTextOutputSelection::LastSequencePosition
38    }
39}
40
41/// Architecture-declared projection from complete logits to one causal-text output.
42#[derive(Debug, Clone, Copy, Eq, PartialEq)]
43#[non_exhaustive]
44pub enum ReplicatedTextOutputSelection {
45    /// Selects the final position on the architecture's sequence axis.
46    LastSequencePosition,
47}
48
49impl ReplicatedTextOutputSelection {
50    /// Returns the mechanical sequence-axis index requested from a backend tensor.
51    pub const fn sequence_index(self) -> i32 {
52        match self {
53            Self::LastSequencePosition => -1,
54        }
55    }
56}
57
58/// Backend implementation route for one source-to-executable weight lowering.
59#[derive(Debug, Clone, Copy, Eq, PartialEq)]
60#[non_exhaustive]
61pub enum WeightLoweringKind {
62    /// The admitted source encoding is retained by the executable operator.
63    Direct,
64    /// An architecture-owned recipe derives the executable tensor from the admitted source.
65    Derived,
66    /// Payload materialization performs an admitted transformation.
67    Transform,
68    /// An architecture recipe derives a tensor that payload materialization then transforms.
69    DerivedTransform,
70}
71
72/// One exact weight lowering implemented by a backend.
73#[derive(Debug, Clone, Eq, PartialEq)]
74pub struct WeightLoweringCapability {
75    /// Exact neutral lowering request implemented by this capability.
76    descriptor: WeightLoweringDescriptor,
77    /// Whether the lowering is direct or transforming.
78    kind: WeightLoweringKind,
79}
80
81impl WeightLoweringCapability {
82    /// Creates one exact backend lowering mechanism.
83    pub fn new(descriptor: WeightLoweringDescriptor, kind: WeightLoweringKind) -> Self {
84        Self { descriptor, kind }
85    }
86
87    /// Returns the admitted source encoding.
88    pub const fn source(&self) -> &SourceTensorEncoding {
89        self.descriptor.source()
90    }
91
92    /// Returns the executable format produced by this mechanism.
93    pub const fn executable(&self) -> LinearFormat {
94        self.descriptor.executable()
95    }
96
97    /// Returns whether materialization retains or transforms the source.
98    pub const fn kind(&self) -> WeightLoweringKind {
99        self.kind
100    }
101
102    /// Returns the exact geometry-bearing lowering request.
103    pub const fn descriptor(&self) -> &WeightLoweringDescriptor {
104        &self.descriptor
105    }
106}
107
108/// Exact source-to-executable lowering query presented to a backend.
109#[derive(Debug, Clone, Eq, PartialEq)]
110pub struct WeightLoweringDescriptor {
111    source: SourceTensorEncoding,
112    executable: LinearFormat,
113    physical_shape: Vec<usize>,
114    logical_shape: Vec<usize>,
115    packed_axis: Option<usize>,
116}
117
118impl WeightLoweringDescriptor {
119    /// Creates a geometry-bearing lowering query.
120    pub fn new(
121        source: SourceTensorEncoding,
122        executable: LinearFormat,
123        physical_shape: Vec<usize>,
124        logical_shape: Vec<usize>,
125        packed_axis: Option<usize>,
126    ) -> Result<Self, ReplicatedTextContractError> {
127        if physical_shape.contains(&0)
128            || logical_shape.contains(&0)
129            || physical_shape.len() != logical_shape.len()
130        {
131            return Err(ReplicatedTextContractError::invalid(
132                "weight lowering requires positive extents and equal physical and logical ranks",
133            ));
134        }
135        if packed_axis.is_some_and(|axis| axis >= logical_shape.len()) {
136            return Err(ReplicatedTextContractError::invalid(
137                "weight lowering packed axis is outside the logical shape",
138            ));
139        }
140        Ok(Self {
141            source,
142            executable,
143            physical_shape,
144            logical_shape,
145            packed_axis,
146        })
147    }
148
149    /// Returns the admitted source encoding.
150    pub const fn source(&self) -> &SourceTensorEncoding {
151        &self.source
152    }
153
154    /// Returns the selected executable format.
155    pub const fn executable(&self) -> LinearFormat {
156        self.executable
157    }
158
159    /// Returns the admitted physical source shape.
160    pub fn physical_shape(&self) -> &[usize] {
161        &self.physical_shape
162    }
163
164    /// Returns the architecture-declared logical shape.
165    pub fn logical_shape(&self) -> &[usize] {
166        &self.logical_shape
167    }
168
169    /// Returns the executable packing axis, when the parameter is packable.
170    pub const fn packed_axis(&self) -> Option<usize> {
171        self.packed_axis
172    }
173
174    /// Returns the exact extent along the packing axis.
175    pub fn packed_extent(&self) -> Option<usize> {
176        self.packed_axis.map(|axis| self.logical_shape[axis])
177    }
178}
179
180/// Weight-residency mechanism implemented by a backend.
181#[derive(Debug, Clone, Copy, Eq, PartialEq)]
182#[non_exhaustive]
183pub enum WeightResidencyMechanism {
184    /// All parameters remain device resident.
185    Resident,
186    /// A bounded device window is staged from host storage.
187    Windowed,
188    /// Bounded host and device windows are populated from disk.
189    DiskStreamed,
190}
191
192/// Physical placement selected for one semantic mutable-state component.
193#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
194#[non_exhaustive]
195pub enum StateComponentPlacement {
196    /// The mutable component remains on the execution device.
197    Device,
198    /// The append-only component is managed by bounded paged storage.
199    Paged,
200}
201
202/// Physical scalar representation selected for native mutable state.
203#[derive(Debug, Clone, Copy, Eq, PartialEq)]
204#[non_exhaustive]
205pub enum StateStorageDtype {
206    /// IEEE half precision.
207    F16,
208    /// Brain floating point.
209    Bf16,
210    /// IEEE single precision.
211    F32,
212    /// IEEE double precision.
213    F64,
214    /// Two IEEE single-precision components.
215    Complex64,
216    /// Signed 32-bit integer.
217    I32,
218    /// Unsigned 32-bit integer.
219    U32,
220}
221
222impl StateStorageDtype {
223    /// Exact bytes occupied by one native state element.
224    pub const fn bytes(self) -> std::num::NonZeroU8 {
225        let bytes = match self {
226            Self::F16 | Self::Bf16 => 2,
227            Self::F32 | Self::I32 | Self::U32 => 4,
228            Self::F64 | Self::Complex64 => 8,
229        };
230        std::num::NonZeroU8::new(bytes).unwrap()
231    }
232
233    /// Whether this representation belongs to the model's floating state family.
234    pub const fn is_floating(self) -> bool {
235        !matches!(self, Self::I32 | Self::U32)
236    }
237
238    /// Resolves an architecture dtype policy without overriding fixed-width tensors.
239    pub const fn resolve(policy: StateTensorDtype, floating: Option<Self>) -> Option<Self> {
240        match policy {
241            StateTensorDtype::Floating => match floating {
242                Some(dtype) if dtype.is_floating() => Some(dtype),
243                _ => None,
244            },
245            StateTensorDtype::Float32 => Some(Self::F32),
246            StateTensorDtype::Int32 => Some(Self::I32),
247            StateTensorDtype::Uint32 => Some(Self::U32),
248        }
249    }
250}
251
252/// Exact state component and placements implemented by a backend mechanism.
253#[derive(Debug, Clone, Eq, PartialEq)]
254pub struct StateComponentMechanism {
255    layer: usize,
256    component: StateComponentPolicy,
257    device_placement: Option<StateComponentPlacement>,
258    paged_placement: Option<StateComponentPlacement>,
259}
260
261impl StateComponentMechanism {
262    /// Describes support for one exact architecture-declared component.
263    pub fn new(
264        layer: usize,
265        component: StateComponentPolicy,
266        device_placement: Option<StateComponentPlacement>,
267        paged_placement: Option<StateComponentPlacement>,
268    ) -> Self {
269        Self {
270            layer,
271            component,
272            device_placement,
273            paged_placement,
274        }
275    }
276
277    /// Returns the architecture-global state layer.
278    pub const fn layer(&self) -> usize {
279        self.layer
280    }
281
282    /// Returns the exact semantic component contract.
283    pub const fn component(&self) -> &StateComponentPolicy {
284        &self.component
285    }
286
287    /// Returns the placement used for a requested state policy.
288    pub const fn placement(
289        &self,
290        policy: &CacheResidencyPolicy,
291    ) -> Option<StateComponentPlacement> {
292        match policy {
293            CacheResidencyPolicy::Device => self.device_placement,
294            CacheResidencyPolicy::Paged(_) => self.paged_placement,
295        }
296    }
297}
298
299/// Exact, family-neutral mutable-state mechanisms reported by a backend.
300#[derive(Debug, Clone, Eq, PartialEq)]
301pub struct StateMechanismCapabilities {
302    floating_state: Option<(TensorDtype, StateStorageDtype)>,
303    components: Vec<StateComponentMechanism>,
304    checkpoint: bool,
305    rollback: bool,
306    reset: bool,
307    prompt_cache: bool,
308    observation_retention: bool,
309}
310
311impl StateMechanismCapabilities {
312    /// Creates a fail-closed report for exact architecture-declared components.
313    pub fn new(components: impl IntoIterator<Item = StateComponentMechanism>) -> Self {
314        Self {
315            floating_state: None,
316            components: components.into_iter().collect(),
317            checkpoint: false,
318            rollback: false,
319            reset: false,
320            prompt_cache: false,
321            observation_retention: false,
322        }
323    }
324
325    /// Binds floating-state support to the exact architecture-selected source dtype.
326    pub fn with_floating_state_dtype(
327        mut self,
328        source: TensorDtype,
329        dtype: StateStorageDtype,
330    ) -> Self {
331        self.floating_state = Some((source, dtype));
332        self
333    }
334
335    /// Returns the source and native representation used for floating-state support queries.
336    pub fn floating_state_dtype(&self) -> Option<(&TensorDtype, StateStorageDtype)> {
337        self.floating_state
338            .as_ref()
339            .map(|(source, dtype)| (source, *dtype))
340    }
341
342    /// Declares transactional checkpoint and rollback facilities.
343    pub const fn with_transactions(mut self, checkpoint: bool, rollback: bool) -> Self {
344        self.checkpoint = checkpoint;
345        self.rollback = rollback;
346        self
347    }
348
349    /// Declares complete state reset support.
350    pub const fn with_reset(mut self, supported: bool) -> Self {
351        self.reset = supported;
352        self
353    }
354
355    /// Declares prompt-cache persistence and restoration support.
356    pub const fn with_prompt_cache(mut self, supported: bool) -> Self {
357        self.prompt_cache = supported;
358        self
359    }
360
361    /// Declares that observed submissions retain every live component.
362    pub const fn with_observation_retention(mut self, supported: bool) -> Self {
363        self.observation_retention = supported;
364        self
365    }
366
367    /// Returns exact supported component mechanisms.
368    pub fn components(&self) -> &[StateComponentMechanism] {
369        &self.components
370    }
371
372    /// Returns whether state checkpoints are implemented.
373    pub const fn checkpoint(&self) -> bool {
374        self.checkpoint
375    }
376
377    /// Returns whether checkpoint rollback is implemented.
378    pub const fn rollback(&self) -> bool {
379        self.rollback
380    }
381
382    /// Returns whether complete reset is implemented.
383    pub const fn reset(&self) -> bool {
384        self.reset
385    }
386
387    /// Returns whether prompt-cache persistence is implemented.
388    pub const fn prompt_cache(&self) -> bool {
389        self.prompt_cache
390    }
391
392    /// Returns whether observation retains every live component.
393    pub const fn observation_retention(&self) -> bool {
394        self.observation_retention
395    }
396}
397
398/// Architecture-valid transform target for one linear parameter.
399#[derive(Debug, Clone, Eq, PartialEq)]
400pub struct ParameterTransformTarget {
401    /// Requested load-time transform.
402    request: QuantizationRequest,
403    /// Executable format produced for this parameter.
404    executable: LinearFormat,
405    /// Exact geometry that the backend lowering must accept.
406    descriptor: WeightLoweringDescriptor,
407}
408
409impl ParameterTransformTarget {
410    /// Creates one architecture-admitted load-time transform target.
411    fn new(
412        request: QuantizationRequest,
413        executable: LinearFormat,
414        descriptor: WeightLoweringDescriptor,
415    ) -> Self {
416        Self {
417            request,
418            executable,
419            descriptor,
420        }
421    }
422
423    /// Returns the caller request selecting this transform.
424    pub const fn request(&self) -> QuantizationRequest {
425        self.request
426    }
427
428    /// Returns the architecture-admitted executable format.
429    pub const fn executable(&self) -> LinearFormat {
430        self.executable
431    }
432
433    /// Returns the exact neutral lowering query.
434    pub const fn descriptor(&self) -> &WeightLoweringDescriptor {
435        &self.descriptor
436    }
437}
438
439/// Architecture declaration of whether and how a parameter may be transformed.
440#[derive(Debug, Clone, Copy, Eq, PartialEq)]
441#[non_exhaustive]
442pub enum ParameterTransformConstraint {
443    /// This parameter is not an executable affine projection weight.
444    None,
445    /// The declared axis is the input/packing axis of a linear parameter.
446    Linear {
447        /// Axis whose extent is grouped or blocked by executable packing.
448        packed_axis: usize,
449    },
450}
451
452/// Architecture-owned semantic role of one logical parameter.
453#[derive(Debug, Clone, Copy, Eq, PartialEq)]
454#[non_exhaustive]
455pub enum ReplicatedTextParameterRole {
456    /// Token lookup table.
457    Embedding,
458    /// Executable affine projection weight.
459    LinearWeight,
460    /// Learned affine projection bias.
461    LinearBias,
462    /// Learned normalization scale or offset.
463    Normalization,
464    /// Physical scale, zero-point, or packed-format companion.
465    FormatCompanion,
466    /// Another architecture-declared non-linear parameter.
467    Other,
468}
469
470/// Architecture-owned location of one replicated-text parameter.
471#[derive(Debug, Clone, Eq, PartialEq)]
472#[non_exhaustive]
473pub enum ReplicatedTextParameterOwner {
474    /// Pinned module selected by a stable architecture role.
475    StaticRole(String),
476    /// One architecture-global execution unit.
477    ExecutionUnit {
478        /// Stable execution-group identity.
479        group: String,
480        /// Group-local architecture-global unit index.
481        unit: usize,
482    },
483}
484
485/// Exact admitted presence or derivation of one logical parameter.
486#[derive(Debug, Clone, Eq, PartialEq)]
487#[non_exhaustive]
488pub enum ReplicatedTextParameterPresence {
489    /// A required physical source was selected.
490    Required,
491    /// An optional physical source was present and selected.
492    OptionalPresent,
493    /// An optional architecture parameter is absent from this artifact.
494    OptionalAbsent,
495    /// The value is tied to another canonical logical parameter.
496    Tied {
497        /// Canonical identity supplying the value.
498        target: String,
499    },
500    /// The value is produced by an architecture-owned recipe.
501    Derived {
502        /// Stable recipe identity.
503        recipe: String,
504    },
505}
506
507impl ReplicatedTextParameterPresence {
508    /// Returns whether selection must choose a backend lowering.
509    pub fn has_physical_source(&self) -> bool {
510        matches!(self, Self::Required | Self::OptionalPresent)
511    }
512}
513
514/// Exact admitted source and executable constraints for one logical parameter.
515#[derive(Debug, Clone, Eq, PartialEq)]
516pub struct ReplicatedTextPhysicalSource {
517    catalog_key: String,
518    tensor: String,
519    shard: PathBuf,
520    output: String,
521    source_encoding: SourceTensorEncoding,
522    encoded_byte_len: u64,
523}
524
525impl ReplicatedTextPhysicalSource {
526    /// Records one exact physical tensor, admitted shard, and selected output.
527    pub fn new(
528        catalog_key: impl Into<String>,
529        tensor: impl Into<String>,
530        shard: impl Into<PathBuf>,
531        output: impl Into<String>,
532        source_encoding: SourceTensorEncoding,
533        encoded_byte_len: u64,
534    ) -> Result<Self, ReplicatedTextContractError> {
535        let catalog_key = catalog_key.into();
536        let tensor = tensor.into();
537        let shard = shard.into();
538        let output = output.into();
539        if catalog_key.trim().is_empty()
540            || tensor.trim().is_empty()
541            || shard.as_os_str().is_empty()
542            || output.trim().is_empty()
543            || encoded_byte_len == 0
544        {
545            return Err(ReplicatedTextContractError::invalid(
546                "physical source key, tensor, shard, output, and byte length must be valid",
547            ));
548        }
549        Ok(Self {
550            catalog_key,
551            tensor,
552            shard,
553            output,
554            source_encoding,
555            encoded_byte_len,
556        })
557    }
558
559    /// Logical key selecting this exact output from the admitted catalog.
560    pub fn catalog_key(&self) -> &str {
561        &self.catalog_key
562    }
563
564    /// Physical tensor identity in the admitted container.
565    pub fn tensor(&self) -> &str {
566        &self.tensor
567    }
568    /// Canonical admitted payload shard.
569    pub fn shard(&self) -> &Path {
570        &self.shard
571    }
572    /// Exact logical output selected from the physical tensor.
573    pub fn output(&self) -> &str {
574        &self.output
575    }
576    /// Exact physical container encoding for this catalog output.
577    pub const fn source_encoding(&self) -> &SourceTensorEncoding {
578        &self.source_encoding
579    }
580    /// Encoded bytes selected for this catalog output.
581    pub const fn encoded_byte_len(&self) -> u64 {
582        self.encoded_byte_len
583    }
584}
585
586/// Exact admitted source and executable constraints for one logical parameter.
587#[derive(Debug, Clone, Eq, PartialEq)]
588pub struct ReplicatedTextParameterRequirement {
589    /// Canonical logical parameter identity.
590    name: String,
591    /// Physical outputs admitted as sources for this logical parameter.
592    sources: Vec<String>,
593    /// Exact shard and multi-output provenance for the physical input.
594    physical_sources: Vec<ReplicatedTextPhysicalSource>,
595    /// All admitted aliases for the logical parameter.
596    aliases: Vec<String>,
597    /// Encoding of the selected physical source, when present.
598    source_encoding: Option<SourceTensorEncoding>,
599    /// Exact selected physical source shape, when present.
600    physical_shape: Option<Vec<usize>>,
601    /// Architecture-declared logical tensor shape.
602    logical_shape: Vec<usize>,
603    /// Architecture-owned semantic parameter role.
604    role: ReplicatedTextParameterRole,
605    /// Architecture-owned static/group/unit location.
606    owner: ReplicatedTextParameterOwner,
607    /// Exact artifact presence, tie, or derivation.
608    presence: ReplicatedTextParameterPresence,
609    /// Architecture-selected native executable format.
610    native_executable: LinearFormat,
611    /// Exact architecture-owned transform eligibility and packing axis.
612    transform: ParameterTransformConstraint,
613    /// Exact encoded-linear primary relationship for a physical companion.
614    linear_companion: Option<(eredu_nn::LinearCompanionRole, String)>,
615    /// Exact architecture output names used when this weight is transformed.
616    transform_companions: Option<(String, String)>,
617    /// Exact stored dtypes which the architecture permits a native companion
618    /// slot to accept during binding.
619    permitted_native_source_dtypes: Vec<eredu_checkpoint::recipe::RecipeDtype>,
620}
621
622impl ReplicatedTextParameterRequirement {
623    /// Creates one exact logical-parameter requirement.
624    #[allow(
625        clippy::too_many_arguments,
626        reason = "the constructor validates one complete immutable catalog record"
627    )]
628    pub fn new(
629        name: impl Into<String>,
630        sources: Vec<String>,
631        physical_sources: Vec<ReplicatedTextPhysicalSource>,
632        aliases: Vec<String>,
633        source_encoding: Option<SourceTensorEncoding>,
634        physical_shape: Option<Vec<usize>>,
635        logical_shape: Vec<usize>,
636        native_executable: LinearFormat,
637        role: ReplicatedTextParameterRole,
638        owner: ReplicatedTextParameterOwner,
639        presence: ReplicatedTextParameterPresence,
640        transform: ParameterTransformConstraint,
641    ) -> Result<Self, ReplicatedTextContractError> {
642        let name = name.into();
643        native_executable
644            .validate()
645            .map_err(|error| ReplicatedTextContractError::invalid(error.to_string()))?;
646        if name.trim().is_empty() {
647            return Err(ReplicatedTextContractError::invalid(
648                "logical parameter identity is empty",
649            ));
650        }
651        if sources.iter().any(|source| source.trim().is_empty())
652            || aliases.iter().any(|alias| alias.trim().is_empty())
653        {
654            return Err(ReplicatedTextContractError::invalid(format!(
655                "logical parameter {name:?} has an empty physical identity"
656            )));
657        }
658        let has_source = !sources.is_empty();
659        let has_physical_facts = source_encoding.is_some() && physical_shape.is_some();
660        if source_encoding.is_some() != physical_shape.is_some()
661            || (has_source && !has_physical_facts)
662        {
663            return Err(ReplicatedTextContractError::invalid(format!(
664                "logical parameter {name:?} has inconsistent source presence"
665            )));
666        }
667        match presence {
668            ReplicatedTextParameterPresence::Required
669            | ReplicatedTextParameterPresence::OptionalPresent
670                if !has_source =>
671            {
672                return Err(ReplicatedTextContractError::invalid(format!(
673                    "physical logical parameter {name:?} has no lowering source"
674                )));
675            }
676            ReplicatedTextParameterPresence::OptionalAbsent
677            | ReplicatedTextParameterPresence::Tied { .. }
678                if has_source =>
679            {
680                return Err(ReplicatedTextContractError::invalid(format!(
681                    "source-free logical parameter {name:?} has a lowering source"
682                )));
683            }
684            _ => {}
685        }
686        let provenance_required =
687            has_source || matches!(presence, ReplicatedTextParameterPresence::Derived { .. });
688        if provenance_required != !physical_sources.is_empty() {
689            return Err(ReplicatedTextContractError::invalid(format!(
690                "logical parameter {name:?} has inconsistent physical provenance"
691            )));
692        }
693        if physical_sources.is_empty() && has_physical_facts {
694            return Err(ReplicatedTextContractError::invalid(format!(
695                "logical parameter {name:?} has physical facts without provenance"
696            )));
697        }
698        if physical_shape
699            .as_ref()
700            .is_some_and(|shape| shape.contains(&0))
701        {
702            return Err(ReplicatedTextContractError::invalid(format!(
703                "logical parameter {name:?} has an invalid physical shape"
704            )));
705        }
706        if logical_shape.contains(&0) {
707            return Err(ReplicatedTextContractError::invalid(format!(
708                "logical parameter {name:?} has an invalid shape {logical_shape:?}"
709            )));
710        }
711        if let ParameterTransformConstraint::Linear { packed_axis } = transform {
712            if packed_axis >= logical_shape.len() {
713                return Err(ReplicatedTextContractError::invalid(format!(
714                    "logical parameter {name:?} has packing axis {packed_axis} outside shape {logical_shape:?}"
715                )));
716            }
717        }
718        let requirement = Self {
719            name,
720            sources,
721            physical_sources,
722            aliases,
723            source_encoding,
724            physical_shape,
725            logical_shape,
726            role,
727            owner,
728            presence,
729            native_executable,
730            transform,
731            linear_companion: None,
732            transform_companions: None,
733            permitted_native_source_dtypes: Vec::new(),
734        };
735        Ok(requirement)
736    }
737
738    /// Explicitly permits exact source dtypes for this architecture parameter.
739    pub fn with_permitted_native_source_dtypes(
740        mut self,
741        dtypes: Vec<eredu_checkpoint::recipe::RecipeDtype>,
742    ) -> Self {
743        self.permitted_native_source_dtypes =
744            dtypes.into_iter().fold(Vec::new(), |mut out, dtype| {
745                if !out.contains(&dtype) {
746                    out.push(dtype);
747                }
748                out
749            });
750        self
751    }
752
753    /// Attaches the exact encoded-linear primary relationship selected by the architecture.
754    pub fn with_linear_companion(
755        mut self,
756        role: eredu_nn::LinearCompanionRole,
757        primary: impl Into<String>,
758    ) -> Result<Self, ReplicatedTextContractError> {
759        let primary = primary.into();
760        if self.role != ReplicatedTextParameterRole::FormatCompanion
761            || primary.trim().is_empty()
762            || primary == self.name
763        {
764            return Err(ReplicatedTextContractError::invalid(format!(
765                "parameter {:?} has an invalid encoded-linear companion relationship",
766                self.name
767            )));
768        }
769        self.linear_companion = Some((role, primary));
770        Ok(self)
771    }
772
773    /// Attaches exact scale and affine-bias output identities for load-time transforms.
774    pub fn with_transform_companions(
775        mut self,
776        scale: impl Into<String>,
777        affine_bias: impl Into<String>,
778    ) -> Result<Self, ReplicatedTextContractError> {
779        let scale = scale.into();
780        let affine_bias = affine_bias.into();
781        if !matches!(self.transform, ParameterTransformConstraint::Linear { .. })
782            || scale.trim().is_empty()
783            || affine_bias.trim().is_empty()
784            || scale == affine_bias
785            || scale == self.name
786            || affine_bias == self.name
787        {
788            return Err(ReplicatedTextContractError::invalid(format!(
789                "parameter {:?} has invalid transform companion identities",
790                self.name
791            )));
792        }
793        self.transform_companions = Some((scale, affine_bias));
794        Ok(self)
795    }
796
797    /// Returns the canonical logical identity.
798    pub fn name(&self) -> &str {
799        &self.name
800    }
801
802    /// Returns exact admitted physical source identities.
803    pub fn sources(&self) -> &[String] {
804        &self.sources
805    }
806
807    /// Returns exact admitted shard and multi-output provenance.
808    pub fn physical_sources(&self) -> &[ReplicatedTextPhysicalSource] {
809        &self.physical_sources
810    }
811
812    /// Returns all architecture-admitted alternative source identities.
813    pub fn aliases(&self) -> &[String] {
814        &self.aliases
815    }
816
817    /// Returns the admitted physical source encoding.
818    pub const fn source_encoding(&self) -> Option<&SourceTensorEncoding> {
819        self.source_encoding.as_ref()
820    }
821
822    /// Returns the selected physical shape, when a source is present.
823    pub fn physical_shape(&self) -> Option<&[usize]> {
824        self.physical_shape.as_deref()
825    }
826
827    /// Returns the architecture-declared logical shape.
828    pub fn logical_shape(&self) -> &[usize] {
829        &self.logical_shape
830    }
831
832    /// Returns the architecture-owned semantic role.
833    pub const fn role(&self) -> ReplicatedTextParameterRole {
834        self.role
835    }
836
837    /// Returns the architecture-owned static/group/unit location.
838    pub const fn owner(&self) -> &ReplicatedTextParameterOwner {
839        &self.owner
840    }
841
842    /// Returns exact artifact presence, tie, or derivation.
843    pub const fn presence(&self) -> &ReplicatedTextParameterPresence {
844        &self.presence
845    }
846
847    /// Returns whether this logical value selects a physical source lowering.
848    ///
849    /// Architecture-derived values may retain a physical lowering source when
850    /// a recipe splits one encoded tensor into several logical parameters.
851    pub fn has_lowering_source(&self) -> bool {
852        !self.sources.is_empty() || !self.physical_sources.is_empty()
853    }
854
855    /// Returns exact transform eligibility and packing geometry.
856    pub const fn transform_constraint(&self) -> ParameterTransformConstraint {
857        self.transform
858    }
859
860    /// Returns the exact encoded-linear primary relationship.
861    pub fn linear_companion(&self) -> Option<(eredu_nn::LinearCompanionRole, &str)> {
862        self.linear_companion
863            .as_ref()
864            .map(|(role, primary)| (*role, primary.as_str()))
865    }
866
867    /// Returns exact scale and affine-bias identities for load-time transforms.
868    pub fn transform_companions(&self) -> Option<(&str, &str)> {
869        self.transform_companions
870            .as_ref()
871            .map(|(scale, bias)| (scale.as_str(), bias.as_str()))
872    }
873
874    /// Returns explicitly permitted native source dtypes.
875    pub fn permitted_native_source_dtypes(&self) -> &[eredu_checkpoint::recipe::RecipeDtype] {
876        &self.permitted_native_source_dtypes
877    }
878
879    /// Returns the architecture-native executable format.
880    pub const fn native_executable(&self) -> LinearFormat {
881        self.native_executable
882    }
883
884    /// Resolves a caller transform through architecture-owned constraints.
885    pub fn transform_target(
886        &self,
887        request: QuantizationRequest,
888    ) -> Result<Option<ParameterTransformTarget>, ReplicatedTextContractError> {
889        let packed_axis = match self.transform {
890            ParameterTransformConstraint::None => return Ok(None),
891            ParameterTransformConstraint::Linear { packed_axis } => packed_axis,
892        };
893        let extent = self.logical_shape[packed_axis];
894        let executable = match request {
895            QuantizationRequest::Affine { group_size, bits } => {
896                let group_size = i32::try_from(group_size).map_err(|_| {
897                    ReplicatedTextContractError::invalid("affine group size exceeds i32")
898                })?;
899                let format = eredu_checkpoint::AffineQuantization::new(group_size, i32::from(bits))
900                    .map_err(|error| ReplicatedTextContractError::invalid(error.to_string()))?;
901                let group_size = usize::try_from(format.group_size).map_err(|_| {
902                    ReplicatedTextContractError::invalid("affine group size is negative")
903                })?;
904                if group_size > extent || !extent.is_multiple_of(group_size) {
905                    return Err(ReplicatedTextContractError::invalid(format!(
906                        "affine group size {group_size} does not divide packed extent {extent}"
907                    )));
908                }
909                LinearFormat::Affine(format)
910            }
911            QuantizationRequest::MxFp4 => {
912                const MXFP4_BLOCK_SIZE: usize = 32;
913                if !extent.is_multiple_of(MXFP4_BLOCK_SIZE) {
914                    return Err(ReplicatedTextContractError::invalid(format!(
915                        "MXFP4 packed extent {extent} is not divisible by block size {MXFP4_BLOCK_SIZE}"
916                    )));
917                }
918                LinearFormat::MxFp4
919            }
920            _ => {
921                return Err(ReplicatedTextContractError::invalid(
922                    "unknown load-time transform request",
923                ))
924            }
925        };
926        let descriptor = self.lowering_descriptor(executable)?;
927        Ok(Some(ParameterTransformTarget::new(
928            request, executable, descriptor,
929        )))
930    }
931
932    /// Forms the exact backend lowering query for one admitted executable format.
933    pub fn lowering_descriptor(
934        &self,
935        executable: LinearFormat,
936    ) -> Result<WeightLoweringDescriptor, ReplicatedTextContractError> {
937        let packed_axis = match self.transform {
938            ParameterTransformConstraint::None => None,
939            ParameterTransformConstraint::Linear { packed_axis } => Some(packed_axis),
940        };
941        let packed_axis = packed_axis
942            .or_else(|| {
943                (self.role == ReplicatedTextParameterRole::Embedding
944                    && executable != LinearFormat::Dense)
945                    .then(|| self.logical_shape.len().checked_sub(1))
946                    .flatten()
947            })
948            .or_else(|| {
949                (matches!(
950                    self.presence,
951                    ReplicatedTextParameterPresence::Derived { .. }
952                ) && executable != LinearFormat::Dense)
953                    .then(|| {
954                        self.physical_shape
955                            .as_ref()
956                            .and_then(|shape| shape.len().checked_sub(1))
957                    })
958                    .flatten()
959            });
960        let alias_backed_packed_output = matches!(
961            self.source_encoding,
962            Some(
963                SourceTensorEncoding::Safetensors(StoredDtype::U32)
964                    | SourceTensorEncoding::RecipeOutput(StoredDtype::U32)
965            )
966        );
967        let lowering_shape = if matches!(
968            self.presence,
969            ReplicatedTextParameterPresence::Derived { .. }
970        ) && !alias_backed_packed_output
971        {
972            self.physical_shape.as_ref().unwrap_or(&self.logical_shape)
973        } else {
974            &self.logical_shape
975        };
976        WeightLoweringDescriptor::new(
977            self.source_encoding.clone().ok_or_else(|| {
978                ReplicatedTextContractError::invalid(format!(
979                    "logical parameter {:?} has no physical lowering source",
980                    self.name
981                ))
982            })?,
983            executable,
984            self.physical_shape.clone().ok_or_else(|| {
985                ReplicatedTextContractError::invalid(format!(
986                    "logical parameter {:?} has no physical source shape",
987                    self.name
988                ))
989            })?,
990            lowering_shape.clone(),
991            packed_axis,
992        )
993    }
994}
995
996/// Invalid public replicated-text contract construction.
997#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
998#[error("invalid replicated text contract: {message}")]
999pub struct ReplicatedTextContractError {
1000    message: String,
1001}
1002
1003/// Static state-access semantics required by an admitted replicated text graph.
1004#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
1005#[non_exhaustive]
1006pub enum ReplicatedTextStateAccess {
1007    /// No mutable token state.
1008    Stateless,
1009    /// Ordinary key/value attention state.
1010    KeyValue,
1011    /// Architecture-declared recurrent or convolutional components only.
1012    Fixed,
1013    /// Key/value attention plus architecture-declared fixed components.
1014    AttentionWithFixed,
1015    /// Compressed-latent attention state without fixed components.
1016    CompressedAttention,
1017    /// Compressed-latent attention plus architecture-declared fixed components.
1018    CompressedAttentionWithFixed,
1019}
1020
1021impl ReplicatedTextContractError {
1022    fn invalid(message: impl Into<String>) -> Self {
1023        Self {
1024            message: message.into(),
1025        }
1026    }
1027
1028    /// Returns the stable semantic diagnostic.
1029    pub fn message(&self) -> &str {
1030        &self.message
1031    }
1032}
1033
1034/// Exact architecture and artifact requirements for replicated text execution.
1035#[derive(Debug, Clone, Eq, PartialEq)]
1036pub struct ReplicatedTextRequirements {
1037    floating_state_source: Option<TensorDtype>,
1038    architecture_identity: String,
1039    /// Optional neural operations required by the architecture equations.
1040    operators: NeuralOperatorCapabilities,
1041    /// Stable architecture-owned execution graph.
1042    execution_graph: ExecutionGraph,
1043    /// Exact group-major execution-unit geometry.
1044    execution_units: ExecutionUnitLayout,
1045    /// Architecture-owned transport semantics in graph-group order.
1046    group_transports: Vec<ArchitectureGroupTransport>,
1047    /// Complete architecture-owned mutable-state geometry.
1048    state_layout: StateLayout,
1049    /// Static state-access semantics used by architecture traversal.
1050    state_access: ReplicatedTextStateAccess,
1051    /// Canonical logical parameter requirements.
1052    parameters: Vec<ReplicatedTextParameterRequirement>,
1053    /// Exact additive prediction/auxiliary parameters selected with this target.
1054    auxiliary_parameters: Vec<ReplicatedTextParameterRequirement>,
1055    derived_recipes: BTreeMap<String, eredu_checkpoint::recipe::DerivedWeightRecipe>,
1056    derived_recipe_outputs: BTreeMap<String, eredu_checkpoint::recipe::RecipeMetadata>,
1057    shared_source_keys: BTreeSet<String>,
1058    grouped_operations: Vec<GroupedOperationRequirement>,
1059}
1060
1061impl ReplicatedTextRequirements {
1062    /// Creates exact requirements from architecture and admitted-artifact facts only.
1063    #[allow(
1064        clippy::too_many_arguments,
1065        reason = "the constructor validates one complete immutable architecture contract"
1066    )]
1067    pub fn new(
1068        architecture_identity: impl Into<String>,
1069        operators: NeuralOperatorCapabilities,
1070        execution_graph: ExecutionGraph,
1071        execution_units: ExecutionUnitLayout,
1072        group_transports: Vec<ArchitectureGroupTransport>,
1073        state_layout: StateLayout,
1074        state_access: ReplicatedTextStateAccess,
1075        parameters: Vec<ReplicatedTextParameterRequirement>,
1076    ) -> Result<Self, ReplicatedTextContractError> {
1077        let architecture_identity = architecture_identity.into();
1078        if architecture_identity.trim().is_empty() {
1079            return Err(ReplicatedTextContractError::invalid(
1080                "architecture identity is empty",
1081            ));
1082        }
1083        if group_transports.len() != execution_graph.groups().len() {
1084            return Err(ReplicatedTextContractError::invalid(format!(
1085                "{} group transports do not match {} execution groups",
1086                group_transports.len(),
1087                execution_graph.groups().len()
1088            )));
1089        }
1090        if execution_units.group_count() != execution_graph.groups().len()
1091            || execution_graph
1092                .groups()
1093                .iter()
1094                .enumerate()
1095                .any(|(index, group)| {
1096                    execution_units
1097                        .group_id(index)
1098                        .is_none_or(|id| id.as_str() != group.id())
1099                })
1100        {
1101            return Err(ReplicatedTextContractError::invalid(
1102                "execution-unit layout group identities differ from the execution graph",
1103            ));
1104        }
1105        validate_state_access_profile(&state_layout, state_access)?;
1106        let mut names = BTreeSet::new();
1107        if parameters
1108            .iter()
1109            .any(|parameter| !names.insert(parameter.name()))
1110        {
1111            return Err(ReplicatedTextContractError::invalid(
1112                "logical parameter identities are not unique",
1113            ));
1114        }
1115        Ok(Self {
1116            floating_state_source: None,
1117            architecture_identity,
1118            operators,
1119            execution_graph,
1120            execution_units,
1121            group_transports,
1122            state_layout,
1123            state_access,
1124            parameters,
1125            auxiliary_parameters: Vec::new(),
1126            derived_recipes: BTreeMap::new(),
1127            derived_recipe_outputs: BTreeMap::new(),
1128            shared_source_keys: BTreeSet::new(),
1129            grouped_operations: Vec::new(),
1130        })
1131    }
1132
1133    /// Records the dtype of the architecture-declared activation source before native selection.
1134    /// This is source metadata, not a request to convert checkpoint weights.
1135    pub fn with_floating_state_source(mut self, dtype: TensorDtype) -> Self {
1136        self.floating_state_source = Some(dtype);
1137        self
1138    }
1139
1140    /// Returns the exact source dtype used to resolve generic floating state.
1141    pub fn floating_state_source(&self) -> Option<&TensorDtype> {
1142        self.floating_state_source.as_ref()
1143    }
1144
1145    /// Attaches exact additive auxiliary parameter requirements before backend selection.
1146    pub fn with_auxiliary_parameters(
1147        mut self,
1148        parameters: Vec<ReplicatedTextParameterRequirement>,
1149        recipes: BTreeMap<String, eredu_checkpoint::recipe::DerivedWeightRecipe>,
1150        outputs: BTreeMap<String, eredu_checkpoint::recipe::RecipeMetadata>,
1151    ) -> Result<Self, ReplicatedTextContractError> {
1152        let primary = self
1153            .parameters
1154            .iter()
1155            .map(|parameter| parameter.name())
1156            .collect::<BTreeSet<_>>();
1157        let mut names = BTreeSet::new();
1158        if parameters
1159            .iter()
1160            .any(|parameter| primary.contains(parameter.name()) || !names.insert(parameter.name()))
1161        {
1162            return Err(ReplicatedTextContractError::invalid(
1163                "auxiliary parameter identities overlap or are not unique",
1164            ));
1165        }
1166        if recipes.keys().ne(outputs.keys())
1167            || recipes
1168                .keys()
1169                .any(|target| !names.contains(target.as_str()))
1170            || recipes
1171                .keys()
1172                .any(|target| self.derived_recipes.contains_key(target))
1173        {
1174            return Err(ReplicatedTextContractError::invalid(
1175                "auxiliary derivations do not match auxiliary parameters",
1176            ));
1177        }
1178        self.derived_recipes.extend(recipes);
1179        self.derived_recipe_outputs.extend(outputs);
1180        self.auxiliary_parameters = parameters;
1181        Ok(self)
1182    }
1183
1184    /// Attaches the exact architecture-owned derivations selected for this artifact.
1185    pub fn with_derived_recipes(
1186        mut self,
1187        recipes: BTreeMap<String, eredu_checkpoint::recipe::DerivedWeightRecipe>,
1188        outputs: BTreeMap<String, eredu_checkpoint::recipe::RecipeMetadata>,
1189    ) -> Result<Self, ReplicatedTextContractError> {
1190        self.set_derived_recipes(recipes, outputs, BTreeSet::new())?;
1191        Ok(self)
1192    }
1193
1194    /// Attaches exact derivations together with architecture-declared physical
1195    /// sources which intentionally feed more than one logical destination.
1196    pub fn with_derived_recipes_and_shared_sources(
1197        mut self,
1198        recipes: BTreeMap<String, eredu_checkpoint::recipe::DerivedWeightRecipe>,
1199        outputs: BTreeMap<String, eredu_checkpoint::recipe::RecipeMetadata>,
1200        shared_source_keys: BTreeSet<String>,
1201    ) -> Result<Self, ReplicatedTextContractError> {
1202        self.set_derived_recipes(recipes, outputs, shared_source_keys)?;
1203        Ok(self)
1204    }
1205
1206    fn set_derived_recipes(
1207        &mut self,
1208        recipes: BTreeMap<String, eredu_checkpoint::recipe::DerivedWeightRecipe>,
1209        outputs: BTreeMap<String, eredu_checkpoint::recipe::RecipeMetadata>,
1210        shared_source_keys: BTreeSet<String>,
1211    ) -> Result<(), ReplicatedTextContractError> {
1212        if recipes.keys().ne(outputs.keys()) {
1213            return Err(ReplicatedTextContractError::invalid(
1214                "derived recipe targets and inferred outputs differ",
1215            ));
1216        }
1217        for source in &shared_source_keys {
1218            let claims = recipes
1219                .values()
1220                .filter(|recipe| recipe.source_keys().contains(&source.as_str()))
1221                .count();
1222            if claims < 2 {
1223                return Err(ReplicatedTextContractError::invalid(format!(
1224                    "declared shared source {source:?} is not claimed by multiple derived targets"
1225                )));
1226            }
1227        }
1228        for target in recipes.keys() {
1229            let recipe = recipes
1230                .get(target)
1231                .expect("recipe target came from the same map");
1232            let parameter = self
1233                .parameters
1234                .iter_mut()
1235                .find(|parameter| parameter.name == *target)
1236                .ok_or_else(|| {
1237                    ReplicatedTextContractError::invalid(format!(
1238                        "derived recipe target {target:?} is not a declared parameter"
1239                    ))
1240                })?;
1241            if matches!(
1242                parameter.presence,
1243                ReplicatedTextParameterPresence::OptionalAbsent
1244                    | ReplicatedTextParameterPresence::Tied { .. }
1245            ) {
1246                return Err(ReplicatedTextContractError::invalid(format!(
1247                    "derived recipe target {target:?} has no independent artifact value"
1248                )));
1249            }
1250            parameter.presence = ReplicatedTextParameterPresence::Derived {
1251                recipe: "architecture.recipe".into(),
1252            };
1253            parameter.sources = recipe
1254                .source_keys()
1255                .into_iter()
1256                .map(str::to_owned)
1257                .collect();
1258        }
1259        self.derived_recipes = recipes;
1260        self.derived_recipe_outputs = outputs;
1261        self.shared_source_keys = shared_source_keys;
1262        Ok(())
1263    }
1264
1265    /// Returns the normalized architecture identity bound during admission.
1266    pub fn architecture_identity(&self) -> &str {
1267        &self.architecture_identity
1268    }
1269
1270    /// Rebinds an admitted physical schema to a new architecture identity which
1271    /// deliberately implements the same complete neutral topology.
1272    pub fn with_extension_architecture_identity(
1273        mut self,
1274        architecture_identity: impl Into<String>,
1275    ) -> Result<Self, ReplicatedTextContractError> {
1276        let architecture_identity = architecture_identity.into();
1277        if architecture_identity.trim().is_empty() {
1278            return Err(ReplicatedTextContractError::invalid(
1279                "extension architecture identity is empty",
1280            ));
1281        }
1282        self.architecture_identity = architecture_identity;
1283        Ok(self)
1284    }
1285
1286    /// Declares exact grouped operations required by this architecture path.
1287    pub fn with_grouped_operations(
1288        mut self,
1289        operations: impl IntoIterator<Item = GroupedOperationRequirement>,
1290    ) -> Self {
1291        self.grouped_operations = operations.into_iter().collect();
1292        self
1293    }
1294
1295    /// Returns required optional neural-operation semantics.
1296    pub const fn operators(&self) -> NeuralOperatorCapabilities {
1297        self.operators
1298    }
1299    /// Returns the architecture-owned execution graph.
1300    pub const fn execution_graph(&self) -> &ExecutionGraph {
1301        &self.execution_graph
1302    }
1303    /// Returns group-major execution-unit geometry.
1304    pub const fn execution_units(&self) -> &ExecutionUnitLayout {
1305        &self.execution_units
1306    }
1307    /// Returns architecture-owned group transports.
1308    pub fn group_transports(&self) -> &[ArchitectureGroupTransport] {
1309        &self.group_transports
1310    }
1311    /// Returns complete mutable-state geometry.
1312    pub const fn state_layout(&self) -> &StateLayout {
1313        &self.state_layout
1314    }
1315    /// Returns the state-access semantics used by typed traversal.
1316    pub const fn state_access(&self) -> ReplicatedTextStateAccess {
1317        self.state_access
1318    }
1319    /// Returns canonical logical parameter requirements.
1320    pub fn parameters(&self) -> &[ReplicatedTextParameterRequirement] {
1321        &self.parameters
1322    }
1323
1324    /// Returns exact additive auxiliary parameter requirements.
1325    pub fn auxiliary_parameters(&self) -> &[ReplicatedTextParameterRequirement] {
1326        &self.auxiliary_parameters
1327    }
1328    /// Returns exact derivations that are part of the selected artifact contract.
1329    pub fn derived_recipes(
1330        &self,
1331    ) -> &BTreeMap<String, eredu_checkpoint::recipe::DerivedWeightRecipe> {
1332        &self.derived_recipes
1333    }
1334    /// Returns admission-time output metadata for every exact derivation.
1335    pub fn derived_recipe_outputs(
1336        &self,
1337    ) -> &BTreeMap<String, eredu_checkpoint::recipe::RecipeMetadata> {
1338        &self.derived_recipe_outputs
1339    }
1340    /// Returns physical sources explicitly declared as shared by the architecture.
1341    pub fn shared_source_keys(&self) -> &BTreeSet<String> {
1342        &self.shared_source_keys
1343    }
1344    /// Returns exact grouped operations required before construction.
1345    pub fn grouped_operations(&self) -> &[GroupedOperationRequirement] {
1346        &self.grouped_operations
1347    }
1348}
1349
1350fn validate_state_access_profile(
1351    layout: &StateLayout,
1352    access: ReplicatedTextStateAccess,
1353) -> Result<(), ReplicatedTextContractError> {
1354    use eredu_core::cache::StateComponentRole;
1355
1356    let roles = (0..layout.len())
1357        .flat_map(|layer| {
1358            layout
1359                .components(layer)
1360                .expect("validated state layout exposes every layer")
1361        })
1362        .map(StateComponentPolicy::role)
1363        .collect::<Vec<_>>();
1364    let ordinary = |role| {
1365        matches!(
1366            role,
1367            StateComponentRole::AttentionKeys | StateComponentRole::AttentionValues
1368        )
1369    };
1370    let compressed = |role| {
1371        matches!(
1372            role,
1373            StateComponentRole::CompressedLatent | StateComponentRole::RotaryKeys
1374        )
1375    };
1376    let fixed = |role| matches!(role, StateComponentRole::Fixed(_));
1377    let has_ordinary = roles.iter().copied().any(ordinary);
1378    let has_compressed = roles.iter().copied().any(compressed);
1379    let has_fixed = roles.iter().copied().any(fixed);
1380    let coherent = match access {
1381        ReplicatedTextStateAccess::Stateless => roles.is_empty(),
1382        ReplicatedTextStateAccess::KeyValue => roles.iter().copied().all(ordinary) && has_ordinary,
1383        ReplicatedTextStateAccess::Fixed => roles.iter().copied().all(fixed) && has_fixed,
1384        ReplicatedTextStateAccess::AttentionWithFixed => {
1385            roles
1386                .iter()
1387                .copied()
1388                .all(|role| ordinary(role) || fixed(role))
1389                && has_ordinary
1390                && has_fixed
1391        }
1392        ReplicatedTextStateAccess::CompressedAttention => {
1393            roles.iter().copied().all(compressed) && has_compressed
1394        }
1395        ReplicatedTextStateAccess::CompressedAttentionWithFixed => {
1396            roles
1397                .iter()
1398                .copied()
1399                .all(|role| compressed(role) || fixed(role))
1400                && has_compressed
1401                && has_fixed
1402        }
1403    };
1404    if !coherent {
1405        return Err(ReplicatedTextContractError::invalid(format!(
1406            "state access profile {access:?} does not match component roles {roles:?}"
1407        )));
1408    }
1409    Ok(())
1410}
1411
1412/// One required grouped-compute mechanism.
1413#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
1414#[non_exhaustive]
1415pub enum GroupedOperationRequirement {
1416    /// Ordinary grouped gated-product output.
1417    GatedProduct,
1418    /// Rank-local gated-product partial with an explicit post-reduce term.
1419    GatedProductTensorParallelPartial,
1420    /// Ordinary grouped ReLU-squared output.
1421    Relu2,
1422    /// Rank-local ReLU-squared partial with an explicit post-reduce term.
1423    Relu2TensorParallelPartial,
1424}
1425
1426/// Generic independently addressable storage facilities implemented by a backend.
1427#[derive(Debug, Clone, Copy, Eq, PartialEq)]
1428pub struct AddressableStorageCapabilities {
1429    bulk_access: bool,
1430    incremental_access: bool,
1431    lease_completion: bool,
1432    maximum_compact_bytes: u64,
1433    tiers: AddressableStorageTiers,
1434}
1435
1436/// Generic storage tiers usable by an independently addressable bank.
1437#[derive(Debug, Clone, Copy, Eq, PartialEq)]
1438pub struct AddressableStorageTiers {
1439    device: bool,
1440    host: bool,
1441    disk: bool,
1442}
1443
1444impl AddressableStorageTiers {
1445    /// Creates one exact tier capability set.
1446    pub const fn new(device: bool, host: bool, disk: bool) -> Self {
1447        Self { device, host, disk }
1448    }
1449
1450    /// Returns whether executable device storage is available.
1451    pub const fn device(self) -> bool {
1452        self.device
1453    }
1454
1455    /// Returns whether host staging storage is available.
1456    pub const fn host(self) -> bool {
1457        self.host
1458    }
1459
1460    /// Returns whether lazy checkpoint-backed storage is available.
1461    pub const fn disk(self) -> bool {
1462        self.disk
1463    }
1464}
1465
1466impl AddressableStorageCapabilities {
1467    /// Creates an exact addressable-storage capability report.
1468    pub const fn new(
1469        bulk_access: bool,
1470        incremental_access: bool,
1471        lease_completion: bool,
1472        maximum_compact_bytes: u64,
1473    ) -> Self {
1474        Self {
1475            bulk_access,
1476            incremental_access,
1477            lease_completion,
1478            maximum_compact_bytes,
1479            tiers: AddressableStorageTiers::new(true, true, true),
1480        }
1481    }
1482
1483    /// Replaces the exact supported storage-tier set.
1484    pub const fn with_tiers(mut self, tiers: AddressableStorageTiers) -> Self {
1485        self.tiers = tiers;
1486        self
1487    }
1488
1489    /// Returns whether bounded multi-row access is implemented.
1490    pub const fn bulk_access(self) -> bool {
1491        self.bulk_access
1492    }
1493
1494    /// Returns whether latency-sensitive incremental access is implemented.
1495    pub const fn incremental_access(self) -> bool {
1496        self.incremental_access
1497    }
1498
1499    /// Returns whether acquisitions remain leased through native completion.
1500    pub const fn lease_completion(self) -> bool {
1501        self.lease_completion
1502    }
1503
1504    /// Returns the largest supported per-acquisition compact bank.
1505    pub const fn maximum_compact_bytes(self) -> u64 {
1506        self.maximum_compact_bytes
1507    }
1508
1509    /// Returns the exact independently addressable storage tiers.
1510    pub const fn tiers(self) -> AddressableStorageTiers {
1511        self.tiers
1512    }
1513}
1514
1515/// Family- and execution-class-neutral backend mechanism report.
1516#[derive(Debug, Clone, Eq, PartialEq)]
1517pub struct BackendMechanismCapabilities {
1518    /// Optional neural operations implemented by the backend.
1519    operators: NeuralOperatorCapabilities,
1520    /// Exact admitted source-to-executable lowerings.
1521    weight_lowerings: Vec<WeightLoweringCapability>,
1522    /// Ordinary parameter residency mechanisms.
1523    weight_residencies: Vec<WeightResidencyMechanism>,
1524    /// Exact mutable-state component and lifecycle mechanisms.
1525    state: StateMechanismCapabilities,
1526    /// Exact session facilities implemented by the constructed session.
1527    session: SessionCapabilities,
1528    /// Prompt-cache persistence mechanism is available.
1529    prompt_cache: bool,
1530    /// Exact completion ownership is implemented for submitted work.
1531    exact_completion: bool,
1532    grouped_operations: Vec<GroupedOperationRequirement>,
1533    indexed_movement: bool,
1534    addressable_storage: Option<AddressableStorageCapabilities>,
1535}
1536
1537impl BackendMechanismCapabilities {
1538    /// Creates a fail-closed mechanism report.
1539    pub fn new(
1540        operators: NeuralOperatorCapabilities,
1541        weight_lowerings: Vec<WeightLoweringCapability>,
1542        weight_residencies: Vec<WeightResidencyMechanism>,
1543        state: StateMechanismCapabilities,
1544    ) -> Self {
1545        Self {
1546            operators,
1547            weight_lowerings,
1548            weight_residencies,
1549            state,
1550            session: SessionCapabilities::default(),
1551            prompt_cache: false,
1552            exact_completion: false,
1553            grouped_operations: Vec::new(),
1554            indexed_movement: false,
1555            addressable_storage: None,
1556        }
1557    }
1558
1559    /// Adds supported session-observation and persistence mechanisms.
1560    pub const fn with_session(mut self, session: SessionCapabilities) -> Self {
1561        self.session = session;
1562        self
1563    }
1564    /// Declares prompt-cache persistence support.
1565    pub const fn with_prompt_cache(mut self, supported: bool) -> Self {
1566        self.prompt_cache = supported;
1567        self
1568    }
1569    /// Declares exact native-completion ownership support.
1570    pub const fn with_exact_completion(mut self, supported: bool) -> Self {
1571        self.exact_completion = supported;
1572        self
1573    }
1574    /// Declares exact grouped operation mechanisms.
1575    pub fn with_grouped_operations(
1576        mut self,
1577        operations: impl IntoIterator<Item = GroupedOperationRequirement>,
1578    ) -> Self {
1579        self.grouped_operations = operations.into_iter().collect();
1580        self
1581    }
1582    /// Declares generic indexed discovery, slicing, remapping, and concatenation.
1583    pub const fn with_indexed_movement(mut self, supported: bool) -> Self {
1584        self.indexed_movement = supported;
1585        self
1586    }
1587    /// Declares generic independently addressable storage facilities.
1588    pub const fn with_addressable_storage(
1589        mut self,
1590        capabilities: AddressableStorageCapabilities,
1591    ) -> Self {
1592        self.addressable_storage = Some(capabilities);
1593        self
1594    }
1595    /// Returns neural-operation mechanisms.
1596    pub const fn operators(&self) -> NeuralOperatorCapabilities {
1597        self.operators
1598    }
1599    /// Returns source-to-executable weight-lowering mechanisms.
1600    pub fn weight_lowerings(&self) -> &[WeightLoweringCapability] {
1601        &self.weight_lowerings
1602    }
1603    /// Returns weight-residency mechanisms.
1604    pub fn weight_residencies(&self) -> &[WeightResidencyMechanism] {
1605        &self.weight_residencies
1606    }
1607    /// Returns exact mutable-state component and lifecycle mechanisms.
1608    pub const fn state(&self) -> &StateMechanismCapabilities {
1609        &self.state
1610    }
1611    /// Returns session-observation and persistence mechanisms.
1612    pub const fn session(&self) -> SessionCapabilities {
1613        self.session
1614    }
1615    /// Returns whether prompt-cache persistence is supported.
1616    pub const fn prompt_cache(&self) -> bool {
1617        self.prompt_cache
1618    }
1619    /// Returns whether exact native-completion ownership is supported.
1620    pub const fn exact_completion(&self) -> bool {
1621        self.exact_completion
1622    }
1623    /// Returns exact grouped operation mechanisms.
1624    pub fn grouped_operations(&self) -> &[GroupedOperationRequirement] {
1625        &self.grouped_operations
1626    }
1627    /// Returns whether generic indexed movement is implemented.
1628    pub const fn indexed_movement(&self) -> bool {
1629        self.indexed_movement
1630    }
1631    /// Returns independently addressable storage facilities, when implemented.
1632    pub const fn addressable_storage(&self) -> Option<AddressableStorageCapabilities> {
1633        self.addressable_storage
1634    }
1635}
1636
1637/// Caller choices resolved while selecting one replicated text realization.
1638#[derive(Debug, Clone, Eq, PartialEq)]
1639pub struct ReplicatedTextSelectionRequest {
1640    max_cached_shards: usize,
1641    /// Requested execution topology.
1642    topology: Option<ParallelTopology>,
1643    /// Requested ordinary parameter residency.
1644    residency: LayerWeightResidency,
1645    /// Requested mutable-state implementation and its exact residency policy.
1646    state: CacheResidencyPolicy,
1647    /// Optional load-time transform.
1648    quantization: Option<QuantizationRequest>,
1649    /// Requested optional session facilities.
1650    session: SessionCapabilities,
1651    /// Whether prompt-cache persistence is requested.
1652    prompt_cache: bool,
1653    /// Whether exact completion ownership is requested.
1654    exact_completion: bool,
1655}
1656
1657impl ReplicatedTextSelectionRequest {
1658    /// Creates a replicated request with fail-closed optional facilities.
1659    pub fn new(residency: LayerWeightResidency, state: CacheResidencyPolicy) -> Self {
1660        Self {
1661            max_cached_shards: residency.max_cached_shards(),
1662            topology: None,
1663            residency,
1664            state,
1665            quantization: None,
1666            session: SessionCapabilities::default(),
1667            prompt_cache: false,
1668            exact_completion: false,
1669        }
1670    }
1671    /// Sets the exact source reader-cache limit independently of weight placement.
1672    pub const fn with_max_cached_shards(mut self, maximum: std::num::NonZeroUsize) -> Self {
1673        self.max_cached_shards = maximum.get();
1674        self
1675    }
1676    /// Returns the exact selected source reader-cache limit.
1677    pub const fn max_cached_shards(&self) -> usize {
1678        self.max_cached_shards
1679    }
1680    /// Sets the requested topology.
1681    pub const fn with_topology(mut self, topology: ParallelTopology) -> Self {
1682        self.topology = Some(topology);
1683        self
1684    }
1685    /// Sets the optional load-time transform.
1686    pub const fn with_quantization(mut self, quantization: QuantizationRequest) -> Self {
1687        self.quantization = Some(quantization);
1688        self
1689    }
1690    /// Sets requested session facilities.
1691    pub const fn with_session(mut self, session: SessionCapabilities) -> Self {
1692        self.session = session;
1693        self
1694    }
1695    /// Requests prompt-cache persistence.
1696    pub const fn with_prompt_cache(mut self, required: bool) -> Self {
1697        self.prompt_cache = required;
1698        self
1699    }
1700    /// Requests exact completion ownership.
1701    pub const fn with_exact_completion(mut self, required: bool) -> Self {
1702        self.exact_completion = required;
1703        self
1704    }
1705    /// Returns the requested topology, where `None` means replicated.
1706    pub const fn topology(&self) -> Option<ParallelTopology> {
1707        self.topology
1708    }
1709    /// Returns the requested weight residency.
1710    pub const fn residency(&self) -> LayerWeightResidency {
1711        self.residency
1712    }
1713    /// Returns the requested state policy.
1714    pub const fn state(&self) -> &CacheResidencyPolicy {
1715        &self.state
1716    }
1717    /// Returns the requested transform.
1718    pub const fn quantization(&self) -> Option<QuantizationRequest> {
1719        self.quantization
1720    }
1721    /// Returns requested session facilities.
1722    pub const fn session(&self) -> SessionCapabilities {
1723        self.session
1724    }
1725    /// Returns whether prompt-cache persistence is requested.
1726    pub const fn prompt_cache(&self) -> bool {
1727        self.prompt_cache
1728    }
1729    /// Returns whether exact completion ownership is requested.
1730    pub const fn exact_completion(&self) -> bool {
1731        self.exact_completion
1732    }
1733}
1734
1735/// Selected lowering for one canonical logical parameter.
1736#[derive(Debug, Clone, Eq, PartialEq)]
1737pub struct SelectedParameterRealization {
1738    /// Canonical logical parameter identity.
1739    name: String,
1740    /// Physical outputs admitted as sources for this logical parameter.
1741    sources: Vec<String>,
1742    physical_sources: Vec<ReplicatedTextPhysicalSource>,
1743    /// Admitted physical encoding.
1744    source_encoding: SourceTensorEncoding,
1745    /// Exact executable format used to construct the architecture module.
1746    executable: LinearFormat,
1747    /// Backend lowering selected for materialization.
1748    lowering: WeightLoweringKind,
1749}
1750
1751/// Exact backend work item for one selected logical parameter.
1752///
1753/// This value joins architecture-owned topology and artifact facts with the
1754/// authoritative selected lowering. A materializer may batch these tasks, but
1755/// it must not replace them with one model-wide transform choice.
1756#[derive(Debug, Clone, Eq, PartialEq)]
1757pub struct ReplicatedTextMaterializationTask {
1758    name: String,
1759    sources: Vec<String>,
1760    physical_sources: Vec<ReplicatedTextPhysicalSource>,
1761    aliases: Vec<String>,
1762    source_encoding: SourceTensorEncoding,
1763    physical_shape: Vec<usize>,
1764    logical_shape: Vec<usize>,
1765    role: ReplicatedTextParameterRole,
1766    owner: ReplicatedTextParameterOwner,
1767    presence: ReplicatedTextParameterPresence,
1768    executable: LinearFormat,
1769    lowering: WeightLoweringKind,
1770    lowering_descriptor: WeightLoweringDescriptor,
1771    derived_recipe: Option<eredu_checkpoint::recipe::DerivedWeightRecipe>,
1772    derived_output: Option<eredu_checkpoint::recipe::RecipeMetadata>,
1773    shared_source_keys: BTreeSet<String>,
1774    permitted_native_source_dtypes: Vec<eredu_checkpoint::recipe::RecipeDtype>,
1775    output_companions: Vec<ReplicatedTextOutputCompanion>,
1776}
1777
1778/// Index-based partition of exact materialization tasks by construction owner.
1779///
1780/// Indices refer to the input task slice used to create the plan. Keeping the
1781/// plan free of borrowed or erased task values lets concrete backends retain
1782/// their own statically dispatched construction path.
1783#[derive(Debug, Clone, Eq, PartialEq)]
1784pub struct ReplicatedTextMaterializationPartitionPlan {
1785    task_count: usize,
1786    static_tasks: Vec<usize>,
1787    unit_tasks: Vec<Vec<usize>>,
1788}
1789
1790impl ReplicatedTextMaterializationPartitionPlan {
1791    /// Returns the number of tasks from which this plan was derived.
1792    pub const fn task_count(&self) -> usize {
1793        self.task_count
1794    }
1795
1796    /// Returns indices owned by architecture-static modules.
1797    pub fn static_task_indices(&self) -> &[usize] {
1798        &self.static_tasks
1799    }
1800
1801    /// Returns task-index partitions in local flattened execution-unit order.
1802    pub fn unit_task_indices(&self) -> &[Vec<usize>] {
1803        &self.unit_tasks
1804    }
1805
1806    /// Borrows static tasks from the exact slice used to construct this plan.
1807    pub fn static_tasks<'a>(
1808        &self,
1809        tasks: &'a [ReplicatedTextMaterializationTask],
1810    ) -> Result<Vec<&'a ReplicatedTextMaterializationTask>, ReplicatedTextContractError> {
1811        self.validate_task_slice(tasks)?;
1812        Ok(self
1813            .static_tasks
1814            .iter()
1815            .map(|index| &tasks[*index])
1816            .collect())
1817    }
1818
1819    /// Borrows unit tasks from the exact slice used to construct this plan.
1820    pub fn unit_tasks<'a>(
1821        &self,
1822        tasks: &'a [ReplicatedTextMaterializationTask],
1823    ) -> Result<Vec<Vec<&'a ReplicatedTextMaterializationTask>>, ReplicatedTextContractError> {
1824        self.validate_task_slice(tasks)?;
1825        Ok(self
1826            .unit_tasks
1827            .iter()
1828            .map(|indices| indices.iter().map(|index| &tasks[*index]).collect())
1829            .collect())
1830    }
1831
1832    fn validate_task_slice(
1833        &self,
1834        tasks: &[ReplicatedTextMaterializationTask],
1835    ) -> Result<(), ReplicatedTextContractError> {
1836        if tasks.len() != self.task_count {
1837            return Err(ReplicatedTextContractError::invalid(format!(
1838                "materialization partition plan expects {} tasks, got {}",
1839                self.task_count,
1840                tasks.len()
1841            )));
1842        }
1843        Ok(())
1844    }
1845}
1846
1847/// One executable-format group of transforming task indices.
1848#[derive(Debug, Clone, Eq, PartialEq)]
1849pub struct ReplicatedTextTransformGroup {
1850    quantization: eredu_checkpoint::WeightQuantization,
1851    task_indices: Vec<usize>,
1852}
1853
1854impl ReplicatedTextTransformGroup {
1855    /// Returns the packed output format shared by every task in this group.
1856    pub const fn quantization(&self) -> eredu_checkpoint::WeightQuantization {
1857        self.quantization
1858    }
1859
1860    /// Returns transforming task indices in original task order.
1861    pub fn task_indices(&self) -> &[usize] {
1862        &self.task_indices
1863    }
1864
1865    /// Borrows this group's tasks from the original task slice.
1866    pub fn tasks<'a>(
1867        &self,
1868        tasks: &'a [ReplicatedTextMaterializationTask],
1869    ) -> Result<Vec<&'a ReplicatedTextMaterializationTask>, ReplicatedTextContractError> {
1870        self.task_indices
1871            .iter()
1872            .map(|index| {
1873                tasks.get(*index).ok_or_else(|| {
1874                    ReplicatedTextContractError::invalid(
1875                        "transform group was applied to a different materialization task slice",
1876                    )
1877                })
1878            })
1879            .collect()
1880    }
1881}
1882
1883/// Architecture-declared output companion for one materialized linear weight.
1884#[derive(Debug, Clone, Eq, PartialEq)]
1885pub struct ReplicatedTextOutputCompanion {
1886    name: String,
1887    role: eredu_nn::LinearCompanionRole,
1888    logical_shape: Vec<usize>,
1889    owner: ParameterGroupOwner,
1890    materialization_task: Option<Box<ReplicatedTextMaterializationTask>>,
1891    catalog_source: Option<ReplicatedTextPhysicalSource>,
1892    derived_recipe: Option<eredu_checkpoint::recipe::DerivedWeightRecipe>,
1893    derived_output: Option<eredu_checkpoint::recipe::RecipeMetadata>,
1894}
1895
1896impl ReplicatedTextOutputCompanion {
1897    /// Creates one exact output companion identity and semantic role.
1898    pub fn new(
1899        name: impl Into<String>,
1900        role: eredu_nn::LinearCompanionRole,
1901        logical_shape: Vec<usize>,
1902        owner: ParameterGroupOwner,
1903    ) -> Result<Self, ReplicatedTextContractError> {
1904        let name = name.into();
1905        if name.trim().is_empty() || logical_shape.is_empty() || logical_shape.contains(&0) {
1906            return Err(ReplicatedTextContractError::invalid(
1907                "materialization output companion identity or geometry is invalid",
1908            ));
1909        }
1910        Ok(Self {
1911            name,
1912            role,
1913            logical_shape,
1914            owner,
1915            materialization_task: None,
1916            catalog_source: None,
1917            derived_recipe: None,
1918            derived_output: None,
1919        })
1920    }
1921
1922    pub(crate) fn with_derived_recipe(
1923        mut self,
1924        recipe: eredu_checkpoint::recipe::DerivedWeightRecipe,
1925        output: eredu_checkpoint::recipe::RecipeMetadata,
1926    ) -> Self {
1927        self.derived_recipe = Some(recipe);
1928        self.derived_output = Some(output);
1929        self
1930    }
1931
1932    /// Returns the exact architecture-declared parameter identity.
1933    pub fn name(&self) -> &str {
1934        &self.name
1935    }
1936
1937    /// Returns the companion's role in the encoded linear parameter.
1938    pub const fn role(&self) -> eredu_nn::LinearCompanionRole {
1939        self.role
1940    }
1941
1942    /// Returns the exact architecture-declared companion geometry.
1943    pub fn logical_shape(&self) -> &[usize] {
1944        &self.logical_shape
1945    }
1946
1947    /// Returns the exact architecture-declared companion owner.
1948    pub const fn owner(&self) -> &ParameterGroupOwner {
1949        &self.owner
1950    }
1951
1952    pub(crate) fn with_materialization_task(
1953        mut self,
1954        task: ReplicatedTextMaterializationTask,
1955    ) -> Result<Self, ReplicatedTextContractError> {
1956        let names_output =
1957            task.name() == self.name || task.aliases().iter().any(|alias| alias == &self.name);
1958        if !names_output || !task.output_companions().is_empty() {
1959            return Err(ReplicatedTextContractError::invalid(format!(
1960                "companion {:?} has an inconsistent standalone materialization task",
1961                self.name
1962            )));
1963        }
1964        self.materialization_task = Some(Box::new(task));
1965        Ok(self)
1966    }
1967
1968    /// Retains exact translated-catalog provenance for a directly loaded companion.
1969    pub fn with_catalog_source(mut self, source: ReplicatedTextPhysicalSource) -> Self {
1970        self.catalog_source = Some(source);
1971        self
1972    }
1973
1974    /// Returns the standalone selected materialization task, when one exists.
1975    ///
1976    /// Generated transform outputs and translated checkpoint catalog outputs
1977    /// instead retain their causal source on the primary task or companion.
1978    pub fn materialization_task(&self) -> Option<&ReplicatedTextMaterializationTask> {
1979        self.materialization_task.as_deref()
1980    }
1981
1982    /// Returns exact translated-catalog provenance for this companion.
1983    pub const fn catalog_source(&self) -> Option<&ReplicatedTextPhysicalSource> {
1984        self.catalog_source.as_ref()
1985    }
1986
1987    /// Returns the architecture-owned companion derivation, when required.
1988    pub const fn derived_recipe(&self) -> Option<&eredu_checkpoint::recipe::DerivedWeightRecipe> {
1989        self.derived_recipe.as_ref()
1990    }
1991
1992    /// Returns admission-time metadata for the derived companion output.
1993    pub const fn derived_output(&self) -> Option<&eredu_checkpoint::recipe::RecipeMetadata> {
1994        self.derived_output.as_ref()
1995    }
1996}
1997
1998impl ReplicatedTextMaterializationTask {
1999    /// Creates one exact, source-backed materialization task selected outside
2000    /// the ordinary replicated-text session lifecycle.
2001    ///
2002    /// This is used by architecture-owned auxiliary modules which share the
2003    /// same physical lowering contract but do not own a text session.
2004    #[allow(clippy::too_many_arguments)]
2005    pub fn from_exact_source(
2006        name: impl Into<String>,
2007        physical_source: ReplicatedTextPhysicalSource,
2008        aliases: Vec<String>,
2009        physical_shape: Vec<usize>,
2010        logical_shape: Vec<usize>,
2011        role: ReplicatedTextParameterRole,
2012        owner: ReplicatedTextParameterOwner,
2013        executable: LinearFormat,
2014        lowering: WeightLoweringKind,
2015        lowering_descriptor: WeightLoweringDescriptor,
2016    ) -> Result<Self, ReplicatedTextContractError> {
2017        let name = name.into();
2018        if name.trim().is_empty()
2019            || physical_shape.is_empty()
2020            || logical_shape.is_empty()
2021            || physical_shape.contains(&0)
2022            || logical_shape.contains(&0)
2023            || lowering_descriptor.source() != physical_source.source_encoding()
2024            || lowering_descriptor.executable() != executable
2025            || lowering_descriptor.physical_shape() != physical_shape
2026            || lowering_descriptor.logical_shape() != logical_shape
2027        {
2028            return Err(ReplicatedTextContractError::invalid(
2029                "exact auxiliary materialization task is internally inconsistent",
2030            ));
2031        }
2032        let source = physical_source.catalog_key().to_owned();
2033        Ok(Self {
2034            name,
2035            sources: vec![source],
2036            physical_sources: vec![physical_source],
2037            aliases,
2038            source_encoding: lowering_descriptor.source().clone(),
2039            physical_shape,
2040            logical_shape,
2041            role,
2042            owner,
2043            presence: ReplicatedTextParameterPresence::Required,
2044            executable,
2045            lowering,
2046            lowering_descriptor,
2047            derived_recipe: None,
2048            derived_output: None,
2049            shared_source_keys: BTreeSet::new(),
2050            permitted_native_source_dtypes: Vec::new(),
2051            output_companions: Vec::new(),
2052        })
2053    }
2054
2055    pub(crate) fn set_output_companions(
2056        &mut self,
2057        mut companions: Vec<ReplicatedTextOutputCompanion>,
2058    ) -> Result<(), ReplicatedTextContractError> {
2059        companions.sort_by(|left, right| {
2060            left.role
2061                .cmp(&right.role)
2062                .then_with(|| left.name.cmp(&right.name))
2063        });
2064        if companions
2065            .windows(2)
2066            .any(|pair| pair[0].name == pair[1].name || pair[0].role == pair[1].role)
2067        {
2068            return Err(ReplicatedTextContractError::invalid(format!(
2069                "materialization task {:?} has duplicate output companions",
2070                self.name
2071            )));
2072        }
2073        let roles = companions
2074            .iter()
2075            .map(|companion| companion.role)
2076            .collect::<Vec<_>>();
2077        let expected = if companions.is_empty()
2078            && matches!(
2079                self.lowering,
2080                WeightLoweringKind::Direct | WeightLoweringKind::Derived
2081            ) {
2082            Vec::new()
2083        } else {
2084            match self.executable {
2085                LinearFormat::Dense | LinearFormat::GgufIQuant { .. } => Vec::new(),
2086                LinearFormat::MxFp4 | LinearFormat::E4M3BlockFp8(_) => {
2087                    vec![eredu_nn::LinearCompanionRole::Scale]
2088                }
2089                LinearFormat::Affine(_) => vec![
2090                    eredu_nn::LinearCompanionRole::Scale,
2091                    eredu_nn::LinearCompanionRole::AffineBias,
2092                ],
2093            }
2094        };
2095        let mut expected = expected;
2096        expected.sort();
2097        if roles != expected {
2098            return Err(ReplicatedTextContractError::invalid(format!(
2099                "materialization task {:?} executable {:?} requires companion roles {:?}, got {:?}",
2100                self.name, self.executable, expected, roles
2101            )));
2102        }
2103        self.output_companions = companions;
2104        Ok(())
2105    }
2106
2107    /// Attaches the architecture's exact selected packed-output companions.
2108    pub fn with_output_companions(
2109        mut self,
2110        companions: Vec<ReplicatedTextOutputCompanion>,
2111    ) -> Result<Self, ReplicatedTextContractError> {
2112        self.set_output_companions(companions)?;
2113        Ok(self)
2114    }
2115
2116    /// Returns the canonical logical parameter identity.
2117    pub fn name(&self) -> &str {
2118        &self.name
2119    }
2120
2121    /// Returns every admitted physical source identity.
2122    pub fn sources(&self) -> &[String] {
2123        &self.sources
2124    }
2125
2126    /// Returns exact shard and translated-output provenance.
2127    pub fn physical_sources(&self) -> &[ReplicatedTextPhysicalSource] {
2128        &self.physical_sources
2129    }
2130
2131    /// Returns every architecture-admitted alias.
2132    pub fn aliases(&self) -> &[String] {
2133        &self.aliases
2134    }
2135
2136    /// Returns architecture-declared physical sources shared by logical outputs.
2137    pub fn shared_source_keys(&self) -> &BTreeSet<String> {
2138        &self.shared_source_keys
2139    }
2140
2141    /// Returns exact source dtypes explicitly permitted for native binding.
2142    pub fn permitted_native_source_dtypes(&self) -> &[eredu_checkpoint::recipe::RecipeDtype] {
2143        &self.permitted_native_source_dtypes
2144    }
2145
2146    /// Returns the exact admitted source encoding.
2147    pub const fn source_encoding(&self) -> &SourceTensorEncoding {
2148        &self.source_encoding
2149    }
2150
2151    /// Returns the admitted physical source geometry.
2152    pub fn physical_shape(&self) -> &[usize] {
2153        &self.physical_shape
2154    }
2155
2156    /// Returns the architecture-declared logical geometry.
2157    pub fn logical_shape(&self) -> &[usize] {
2158        &self.logical_shape
2159    }
2160
2161    /// Returns the architecture-owned semantic parameter role.
2162    pub const fn role(&self) -> ReplicatedTextParameterRole {
2163        self.role
2164    }
2165
2166    /// Returns the architecture-owned module location.
2167    pub const fn owner(&self) -> &ReplicatedTextParameterOwner {
2168        &self.owner
2169    }
2170
2171    /// Returns the exact admitted presence or derivation.
2172    pub const fn presence(&self) -> &ReplicatedTextParameterPresence {
2173        &self.presence
2174    }
2175
2176    /// Returns the selected executable format.
2177    pub const fn executable(&self) -> LinearFormat {
2178        self.executable
2179    }
2180
2181    /// Returns the selected backend lowering mechanism.
2182    pub const fn lowering(&self) -> WeightLoweringKind {
2183        self.lowering
2184    }
2185
2186    /// Returns the complete geometry-bearing lowering request.
2187    pub const fn lowering_descriptor(&self) -> &WeightLoweringDescriptor {
2188        &self.lowering_descriptor
2189    }
2190
2191    /// Returns the architecture-owned derivation, when this output is derived.
2192    pub const fn derived_recipe(&self) -> Option<&eredu_checkpoint::recipe::DerivedWeightRecipe> {
2193        self.derived_recipe.as_ref()
2194    }
2195
2196    /// Returns admission-time metadata for the derived output.
2197    pub const fn derived_output(&self) -> Option<&eredu_checkpoint::recipe::RecipeMetadata> {
2198        self.derived_output.as_ref()
2199    }
2200
2201    /// Returns exact output companion identities declared by the architecture.
2202    pub fn output_companions(&self) -> &[ReplicatedTextOutputCompanion] {
2203        &self.output_companions
2204    }
2205
2206    /// Returns the exact source recipe selected for this task.
2207    ///
2208    /// Direct tasks are represented as a full selection of their single
2209    /// admitted source. Derived tasks return the architecture-owned recipe
2210    /// without reconstructing it from a checkpoint catalog.
2211    pub fn source_recipe(
2212        &self,
2213    ) -> Result<eredu_checkpoint::recipe::DerivedWeightRecipe, ReplicatedTextContractError> {
2214        let expects_recipe = matches!(
2215            self.lowering,
2216            WeightLoweringKind::Derived | WeightLoweringKind::DerivedTransform
2217        );
2218        match (expects_recipe, self.derived_recipe.as_ref()) {
2219            (true, Some(recipe)) => {
2220                let declared = self
2221                    .sources
2222                    .iter()
2223                    .map(String::as_str)
2224                    .collect::<BTreeSet<_>>();
2225                let consumed = recipe.source_keys().into_iter().collect::<BTreeSet<_>>();
2226                if declared != consumed {
2227                    return Err(ReplicatedTextContractError::invalid(format!(
2228                        "materialization task {:?} recipe sources differ from its exact source catalog",
2229                        self.name
2230                    )));
2231                }
2232                Ok(recipe.clone())
2233            }
2234            (false, None) => {
2235                let [source] = self.sources.as_slice() else {
2236                    return Err(ReplicatedTextContractError::invalid(format!(
2237                        "direct materialization task {:?} must name exactly one source",
2238                        self.name
2239                    )));
2240                };
2241                Ok(eredu_checkpoint::recipe::DerivedWeightRecipe::source(
2242                    source.clone(),
2243                    eredu_checkpoint::store::TensorSelection::Full,
2244                ))
2245            }
2246            (true, None) => Err(ReplicatedTextContractError::invalid(format!(
2247                "derived materialization task {:?} has no exact recipe",
2248                self.name
2249            ))),
2250            (false, Some(_)) => Err(ReplicatedTextContractError::invalid(format!(
2251                "direct materialization task {:?} unexpectedly carries a recipe",
2252                self.name
2253            ))),
2254        }
2255    }
2256}
2257
2258/// Partitions exact tasks against the architecture-global execution layout.
2259pub fn plan_replicated_text_materialization_tasks(
2260    tasks: &[ReplicatedTextMaterializationTask],
2261    layout: &ExecutionUnitLayout,
2262) -> Result<ReplicatedTextMaterializationPartitionPlan, ReplicatedTextContractError> {
2263    let mut static_tasks = Vec::new();
2264    let mut unit_tasks = vec![Vec::new(); layout.len()];
2265    for (task_index, task) in tasks.iter().enumerate() {
2266        match task.owner() {
2267            ReplicatedTextParameterOwner::StaticRole(_) => static_tasks.push(task_index),
2268            ReplicatedTextParameterOwner::ExecutionUnit { group, unit } => {
2269                let group_index = (0..layout.group_count())
2270                    .find(|index| {
2271                        layout
2272                            .group_id(*index)
2273                            .is_some_and(|id| id.as_str() == group)
2274                    })
2275                    .ok_or_else(|| {
2276                        ReplicatedTextContractError::invalid(format!(
2277                            "exact task {:?} names unknown execution group {group:?}",
2278                            task.name()
2279                        ))
2280                    })?;
2281                let ordinal = layout.ordinal(group_index, *unit).ok_or_else(|| {
2282                    ReplicatedTextContractError::invalid(format!(
2283                        "exact task {:?} names unknown unit {unit} in group {group:?}",
2284                        task.name()
2285                    ))
2286                })?;
2287                unit_tasks[ordinal].push(task_index);
2288            }
2289        }
2290    }
2291    Ok(ReplicatedTextMaterializationPartitionPlan {
2292        task_count: tasks.len(),
2293        static_tasks,
2294        unit_tasks,
2295    })
2296}
2297
2298/// Partitions exact tasks into one retained rank-local execution-unit order.
2299pub fn plan_local_replicated_text_materialization_tasks(
2300    tasks: &[ReplicatedTextMaterializationTask],
2301    global_layout: &ExecutionUnitLayout,
2302    addresses: &[crate::ExecutionUnitAddress],
2303) -> Result<ReplicatedTextMaterializationPartitionPlan, ReplicatedTextContractError> {
2304    if addresses.is_empty() {
2305        return Err(ReplicatedTextContractError::invalid(
2306            "local partition has no selected execution units",
2307        ));
2308    }
2309    let mut seen = BTreeSet::new();
2310    for address in addresses {
2311        if global_layout.address(
2312            global_layout
2313                .ordinal(address.group(), address.index())
2314                .unwrap_or(usize::MAX),
2315        ) != Some(*address)
2316        {
2317            return Err(ReplicatedTextContractError::invalid(format!(
2318                "local partition names unknown global unit {}.{}",
2319                address.group(),
2320                address.index()
2321            )));
2322        }
2323        if !seen.insert((address.group(), address.index())) {
2324            return Err(ReplicatedTextContractError::invalid(format!(
2325                "local partition repeats global unit {}.{}",
2326                address.group(),
2327                address.index()
2328            )));
2329        }
2330    }
2331
2332    let mut static_tasks = Vec::new();
2333    let mut unit_tasks = vec![Vec::new(); addresses.len()];
2334    for (task_index, task) in tasks.iter().enumerate() {
2335        match task.owner() {
2336            ReplicatedTextParameterOwner::StaticRole(_) => static_tasks.push(task_index),
2337            ReplicatedTextParameterOwner::ExecutionUnit { group, unit } => {
2338                let local = addresses
2339                    .iter()
2340                    .position(|address| {
2341                        global_layout
2342                            .group_id(address.group())
2343                            .is_some_and(|id| id.as_str() == group)
2344                            && address.index() == *unit
2345                    })
2346                    .ok_or_else(|| {
2347                        ReplicatedTextContractError::invalid(format!(
2348                            "local task {:?} has no owned global unit {group}.{unit}",
2349                            task.name()
2350                        ))
2351                    })?;
2352                unit_tasks[local].push(task_index);
2353            }
2354        }
2355    }
2356    Ok(ReplicatedTextMaterializationPartitionPlan {
2357        task_count: tasks.len(),
2358        static_tasks,
2359        unit_tasks,
2360    })
2361}
2362
2363/// Returns every primary and companion produced by local transformation.
2364pub fn locally_materialized_replicated_text_outputs(
2365    tasks: &[ReplicatedTextMaterializationTask],
2366) -> BTreeSet<String> {
2367    tasks
2368        .iter()
2369        .filter(|task| {
2370            matches!(
2371                task.lowering(),
2372                WeightLoweringKind::Transform | WeightLoweringKind::DerivedTransform
2373            )
2374        })
2375        .flat_map(|task| {
2376            std::iter::once(task.name().to_owned()).chain(
2377                task.output_companions()
2378                    .iter()
2379                    .map(|companion| companion.name().to_owned()),
2380            )
2381        })
2382        .collect()
2383}
2384
2385/// Groups transforming tasks by exact packed output format.
2386///
2387/// Groups and indices preserve first-observed task order. Direct tasks are not
2388/// included, and a transforming task without a packed format fails closed.
2389pub fn group_replicated_text_transform_tasks(
2390    tasks: &[ReplicatedTextMaterializationTask],
2391) -> Result<Vec<ReplicatedTextTransformGroup>, ReplicatedTextContractError> {
2392    let mut groups = Vec::<ReplicatedTextTransformGroup>::new();
2393    for (task_index, task) in tasks.iter().enumerate().filter(|(_, task)| {
2394        matches!(
2395            task.lowering(),
2396            WeightLoweringKind::Transform | WeightLoweringKind::DerivedTransform
2397        )
2398    }) {
2399        let quantization = task.executable().weight_quantization().ok_or_else(|| {
2400            ReplicatedTextContractError::invalid(format!(
2401                "selected materialization task {:?} has no packed output format",
2402                task.name()
2403            ))
2404        })?;
2405        if let Some(group) = groups
2406            .iter_mut()
2407            .find(|group| group.quantization == quantization)
2408        {
2409            group.task_indices.push(task_index);
2410        } else {
2411            groups.push(ReplicatedTextTransformGroup {
2412                quantization,
2413                task_indices: vec![task_index],
2414            });
2415        }
2416    }
2417    Ok(groups)
2418}
2419
2420/// Computes the exact executable storage charged to one selected task.
2421///
2422/// Direct tasks retain their admitted physical or derived output bytes.
2423/// Transform tasks replace those bytes with the packed weight and exactly the
2424/// companion roles selected by the executable format.
2425pub fn selected_materialization_task_bytes(
2426    task: &ReplicatedTextMaterializationTask,
2427) -> Result<u64, ReplicatedTextContractError> {
2428    let transforms = matches!(
2429        task.lowering(),
2430        WeightLoweringKind::Transform | WeightLoweringKind::DerivedTransform
2431    );
2432    if !transforms {
2433        if let Some(output) = task.derived_output() {
2434            return Ok(output.byte_len());
2435        }
2436        return task
2437            .physical_sources()
2438            .iter()
2439            .try_fold(0u64, |total, source| {
2440                total.checked_add(source.encoded_byte_len()).ok_or_else(|| {
2441                    ReplicatedTextContractError::invalid(format!(
2442                        "materialization task {:?} physical byte total overflowed",
2443                        task.name()
2444                    ))
2445                })
2446            });
2447    }
2448
2449    let dtype = task
2450        .source_encoding()
2451        .scalar_dtype()
2452        .map(eredu_checkpoint::recipe::RecipeDtype::from)
2453        .ok_or_else(|| {
2454            ReplicatedTextContractError::invalid(format!(
2455                "materialization task {:?} transforms a non-scalar source",
2456                task.name()
2457            ))
2458        })?;
2459    let source_bytes = task
2460        .derived_output()
2461        .map(|output| output.byte_len())
2462        .or_else(|| {
2463            task.physical_sources()
2464                .first()
2465                .map(|source| source.encoded_byte_len())
2466        })
2467        .ok_or_else(|| {
2468            ReplicatedTextContractError::invalid(format!(
2469                "materialization task {:?} has no source byte extent",
2470                task.name()
2471            ))
2472        })?;
2473    let metadata = eredu_checkpoint::recipe::RecipeMetadata {
2474        shape: task.logical_shape().to_vec(),
2475        dtype,
2476        byte_len: source_bytes,
2477    };
2478    crate::selected_addressable_parameter_bytes(task, &metadata)
2479        .map_err(|error| ReplicatedTextContractError::invalid(error.to_string()))
2480}
2481
2482/// Projects an authoritative selection into exact materialization work.
2483///
2484/// Every selected parameter must agree with its immutable requirement. The
2485/// returned sequence preserves selected-parameter order and contains no
2486/// model-wide quantization or transform value.
2487pub fn replicated_text_materialization_tasks(
2488    selected: &SelectedReplicatedTextRealization,
2489) -> Result<Vec<ReplicatedTextMaterializationTask>, ReplicatedTextContractError> {
2490    if selected.materialization_tasks.is_empty() && !selected.parameters.is_empty() {
2491        return Err(ReplicatedTextContractError::invalid(
2492            "selected realization omitted its authoritative materialization tasks",
2493        ));
2494    }
2495    Ok(selected.materialization_tasks.clone())
2496}
2497
2498fn build_replicated_text_materialization_tasks(
2499    selected: &SelectedReplicatedTextRealization,
2500) -> Result<Vec<ReplicatedTextMaterializationTask>, ReplicatedTextContractError> {
2501    build_materialization_tasks(
2502        selected.requirements(),
2503        selected.requirements().parameters(),
2504        selected.parameters(),
2505    )
2506}
2507
2508fn build_materialization_tasks(
2509    requirements: &ReplicatedTextRequirements,
2510    parameter_requirements: &[ReplicatedTextParameterRequirement],
2511    selected_parameters: &[SelectedParameterRealization],
2512) -> Result<Vec<ReplicatedTextMaterializationTask>, ReplicatedTextContractError> {
2513    let mut tasks = selected_parameters
2514        .iter()
2515        .map(|realization| {
2516            let requirement = parameter_requirements
2517                .iter()
2518                .find(|requirement| requirement.name() == realization.name())
2519                .ok_or_else(|| {
2520                    ReplicatedTextContractError::invalid(format!(
2521                        "selected parameter {:?} has no architecture requirement",
2522                        realization.name()
2523                    ))
2524                })?;
2525            if requirement.sources() != realization.sources()
2526                || requirement.physical_sources() != realization.physical_sources()
2527                || requirement.source_encoding() != Some(realization.source_encoding())
2528            {
2529                return Err(ReplicatedTextContractError::invalid(format!(
2530                    "selected parameter {:?} changed admitted source provenance",
2531                    realization.name()
2532                )));
2533            }
2534            let physical_shape = requirement.physical_shape().ok_or_else(|| {
2535                ReplicatedTextContractError::invalid(format!(
2536                    "selected parameter {:?} has no physical geometry",
2537                    realization.name()
2538                ))
2539            })?;
2540            let lowering_descriptor = requirement.lowering_descriptor(realization.executable())?;
2541            if lowering_descriptor.source() != realization.source_encoding() {
2542                return Err(ReplicatedTextContractError::invalid(format!(
2543                    "selected parameter {:?} changed its lowering source encoding",
2544                    realization.name()
2545                )));
2546            }
2547            let mut derived_recipe = requirements
2548                .derived_recipes()
2549                .get(realization.name())
2550                .cloned();
2551            let derived_output = requirements
2552                .derived_recipe_outputs()
2553                .get(realization.name())
2554                .cloned();
2555            if derived_recipe.is_some() != derived_output.is_some() {
2556                return Err(ReplicatedTextContractError::invalid(format!(
2557                    "selected parameter {:?} has incomplete derived metadata",
2558                    realization.name()
2559                )));
2560            }
2561            if derived_recipe.is_none()
2562                && matches!(
2563                    realization.lowering(),
2564                    WeightLoweringKind::Derived | WeightLoweringKind::DerivedTransform
2565                )
2566            {
2567                let [source] = realization.sources() else {
2568                    return Err(ReplicatedTextContractError::invalid(format!(
2569                        "derived selected parameter {:?} has no exact recipe and does not name one source",
2570                        realization.name()
2571                    )));
2572                };
2573                derived_recipe = Some(
2574                    eredu_checkpoint::recipe::DerivedWeightRecipe::source(
2575                        source.clone(),
2576                        eredu_checkpoint::store::TensorSelection::Full,
2577                    ),
2578                );
2579            }
2580            Ok(ReplicatedTextMaterializationTask {
2581                name: realization.name().to_owned(),
2582                sources: realization.sources().to_vec(),
2583                physical_sources: realization.physical_sources().to_vec(),
2584                aliases: requirement.aliases().to_vec(),
2585                source_encoding: realization.source_encoding().clone(),
2586                physical_shape: physical_shape.to_vec(),
2587                logical_shape: requirement.logical_shape().to_vec(),
2588                role: requirement.role(),
2589                owner: requirement.owner().clone(),
2590                presence: requirement.presence().clone(),
2591                executable: realization.executable(),
2592                lowering: realization.lowering(),
2593                lowering_descriptor,
2594                derived_recipe,
2595                derived_output,
2596                shared_source_keys: requirements.shared_source_keys().clone(),
2597                permitted_native_source_dtypes: requirement
2598                    .permitted_native_source_dtypes()
2599                    .to_vec(),
2600                output_companions: Vec::new(),
2601            })
2602        })
2603        .collect::<Result<Vec<_>, _>>()?;
2604
2605    let task_by_name = tasks
2606        .iter()
2607        .map(|task| (task.name().to_owned(), task.clone()))
2608        .collect::<BTreeMap<_, _>>();
2609    let requirement_by_name = parameter_requirements
2610        .iter()
2611        .map(|requirement| (requirement.name(), requirement))
2612        .collect::<BTreeMap<_, _>>();
2613    let mut declared = BTreeMap::<String, Vec<ReplicatedTextOutputCompanion>>::new();
2614    let mut companion_names = BTreeSet::new();
2615    for requirement in parameter_requirements {
2616        let Some((role, primary)) = requirement.linear_companion() else {
2617            continue;
2618        };
2619        let owner = parameter_group_owner(requirement.owner())?;
2620        let mut companion = ReplicatedTextOutputCompanion::new(
2621            requirement.name(),
2622            role,
2623            requirement.logical_shape().to_vec(),
2624            owner,
2625        )?;
2626        if let Some(task) = task_by_name.get(requirement.name()) {
2627            companion = companion.with_materialization_task(task.clone())?;
2628        }
2629        declared
2630            .entry(primary.to_owned())
2631            .or_default()
2632            .push(companion);
2633        companion_names.insert(requirement.name().to_owned());
2634    }
2635    for task in &mut tasks {
2636        let transforms = matches!(
2637            task.lowering(),
2638            WeightLoweringKind::Transform | WeightLoweringKind::DerivedTransform
2639        );
2640        let outputs = if transforms {
2641            let requirement = requirement_by_name.get(task.name()).ok_or_else(|| {
2642                ReplicatedTextContractError::invalid(format!(
2643                    "selected task {:?} has no retained requirement",
2644                    task.name()
2645                ))
2646            })?;
2647            let (scale, affine_bias) = requirement.transform_companions().ok_or_else(|| {
2648                ReplicatedTextContractError::invalid(format!(
2649                    "transformed task {:?} has no architecture-selected companion identities",
2650                    task.name()
2651                ))
2652            })?;
2653            let quantization = task.executable().weight_quantization().ok_or_else(|| {
2654                ReplicatedTextContractError::invalid(format!(
2655                    "transformed task {:?} selected a non-quantized executable",
2656                    task.name()
2657                ))
2658            })?;
2659            let mut shape = task.logical_shape().to_vec();
2660            let input = shape.last_mut().ok_or_else(|| {
2661                ReplicatedTextContractError::invalid("transformed scalar parameter")
2662            })?;
2663            let group = usize::try_from(quantization.group_size()).map_err(|_| {
2664                ReplicatedTextContractError::invalid("transform group size exceeds usize")
2665            })?;
2666            if group == 0 || !input.is_multiple_of(group) {
2667                return Err(ReplicatedTextContractError::invalid(format!(
2668                    "transformed task {:?} has incompatible companion geometry",
2669                    task.name()
2670                )));
2671            }
2672            *input /= group;
2673            let owner = parameter_group_owner(task.owner())?;
2674            let mut outputs = vec![ReplicatedTextOutputCompanion::new(
2675                scale,
2676                eredu_nn::LinearCompanionRole::Scale,
2677                shape.clone(),
2678                owner.clone(),
2679            )?];
2680            if quantization.has_biases() {
2681                outputs.push(ReplicatedTextOutputCompanion::new(
2682                    affine_bias,
2683                    eredu_nn::LinearCompanionRole::AffineBias,
2684                    shape,
2685                    owner,
2686                )?);
2687            }
2688            outputs
2689        } else {
2690            declared.remove(task.name()).unwrap_or_default()
2691        };
2692        task.set_output_companions(outputs)?;
2693    }
2694    if !declared.is_empty() {
2695        return Err(ReplicatedTextContractError::invalid(format!(
2696            "selected companion primaries have no materialization task: {:?}",
2697            declared.keys().collect::<Vec<_>>()
2698        )));
2699    }
2700    tasks.retain(|task| !companion_names.contains(task.name()));
2701    Ok(tasks)
2702}
2703
2704fn parameter_group_owner(
2705    owner: &ReplicatedTextParameterOwner,
2706) -> Result<ParameterGroupOwner, ReplicatedTextContractError> {
2707    match owner {
2708        ReplicatedTextParameterOwner::StaticRole(role) => {
2709            Ok(ParameterGroupOwner::static_role(role.clone()))
2710        }
2711        ReplicatedTextParameterOwner::ExecutionUnit { group, unit } => {
2712            let group = ExecutionGroupId::new(group.clone())
2713                .map_err(|error| ReplicatedTextContractError::invalid(error.to_string()))?;
2714            Ok(ParameterGroupOwner::execution_unit(group, *unit))
2715        }
2716    }
2717}
2718
2719/// Projects selected text materialization into one exact architecture partition.
2720///
2721/// Encoded-linear companions are reconstructed from the architecture's
2722/// validated physical parameter groups and remain atomic with their primary
2723/// task. If a partition would own only part of such a physical family, the
2724/// complete projection fails instead of retaining an unowned output.
2725pub fn partitioned_replicated_text_materialization_tasks<G, A>(
2726    selected: &SelectedReplicatedTextRealization,
2727    parameters: &ArchitectureParameterDescription,
2728    partition: &ArchitecturePartition<G, A>,
2729) -> Result<Vec<ReplicatedTextMaterializationTask>, ReplicatedTextContractError> {
2730    let tasks = replicated_text_materialization_tasks(selected)?;
2731    partition_selected_replicated_text_materialization_tasks(&tasks, parameters, partition)
2732}
2733
2734/// Completes rank projection for physical tasks selected before backend resources exist.
2735///
2736/// This attaches architecture-declared atomic companions and removes tasks not owned by
2737/// the exact partition. It never reselects a source, encoding, executable format, recipe,
2738/// or lowering from the architecture topology.
2739pub fn partition_selected_replicated_text_materialization_tasks<G, A>(
2740    tasks: &[ReplicatedTextMaterializationTask],
2741    parameters: &ArchitectureParameterDescription,
2742    partition: &ArchitecturePartition<G, A>,
2743) -> Result<Vec<ReplicatedTextMaterializationTask>, ReplicatedTextContractError> {
2744    let mut tasks = tasks.to_vec();
2745    let mut companions = BTreeMap::<String, Vec<ReplicatedTextOutputCompanion>>::new();
2746    let mut all_targets = BTreeSet::new();
2747    let mut owned_targets = BTreeSet::new();
2748    for tagged in parameters.groups() {
2749        let local = partition.parameter_bindings().iter().any(|binding| {
2750            binding.owner() == tagged.owner()
2751                && parameter_groups_have_same_members(binding.group(), tagged.group())
2752        });
2753        let group_targets = tagged
2754            .members()
2755            .iter()
2756            .map(|member| member.target())
2757            .collect::<BTreeSet<_>>();
2758        for member in tagged.members() {
2759            if !all_targets.insert(member.target().to_owned()) {
2760                return Err(ReplicatedTextContractError::invalid(format!(
2761                    "architecture parameter target {:?} appears more than once",
2762                    member.target()
2763                )));
2764            }
2765            if local {
2766                owned_targets.insert(member.target().to_owned());
2767            }
2768            match (member.linear_companion(), member.linear_companion_of()) {
2769                (None, None) => {}
2770                (Some(role), Some(primary)) if group_targets.contains(primary) && local => {
2771                    companions.entry(primary.to_owned()).or_default().push(
2772                        ReplicatedTextOutputCompanion::new(
2773                            member.target(),
2774                            role,
2775                            member.global_shape().to_vec(),
2776                            tagged.owner().clone(),
2777                        )?,
2778                    );
2779                }
2780                (Some(_), Some(primary)) if group_targets.contains(primary) => {}
2781                (Some(_), Some(primary)) => {
2782                    return Err(ReplicatedTextContractError::invalid(format!(
2783                        "physical companion {:?} names primary {primary:?} outside its atomic parameter group",
2784                        member.target()
2785                    )));
2786                }
2787                _ => {
2788                    return Err(ReplicatedTextContractError::invalid(format!(
2789                        "physical parameter {:?} has incomplete companion metadata",
2790                        member.target()
2791                    )));
2792                }
2793            }
2794        }
2795    }
2796
2797    let mut topology_targets = BTreeMap::<String, String>::new();
2798    let mut target_claims = BTreeMap::<String, String>::new();
2799    for task in &tasks {
2800        let matches = std::iter::once(task.name())
2801            .chain(task.aliases().iter().map(String::as_str))
2802            .filter(|candidate| all_targets.contains(*candidate))
2803            .collect::<BTreeSet<_>>();
2804        if matches.len() != 1 {
2805            return Err(ReplicatedTextContractError::invalid(format!(
2806                "selected materialization output {:?} resolves to {} architecture topology targets through its canonical identity and admitted aliases: {:?}",
2807                task.name(),
2808                matches.len(),
2809                matches
2810            )));
2811        }
2812        let target = matches.first().expect("one topology target was validated");
2813        if let Some(previous) = target_claims.insert((*target).to_owned(), task.name().to_owned()) {
2814            return Err(ReplicatedTextContractError::invalid(format!(
2815                "selected materialization outputs {previous:?} and {:?} ambiguously resolve to architecture target {target:?}",
2816                task.name()
2817            )));
2818        }
2819        topology_targets.insert(task.name().to_owned(), (*target).to_owned());
2820    }
2821    for task in &mut tasks {
2822        let topology_target = topology_targets
2823            .get(task.name())
2824            .expect("every task has one validated topology target");
2825        if !owned_targets.contains(topology_target) {
2826            continue;
2827        }
2828        let mut declared = companions.remove(topology_target).unwrap_or_default();
2829        declared.sort_by(|left, right| {
2830            left.role()
2831                .cmp(&right.role())
2832                .then_with(|| left.name().cmp(right.name()))
2833        });
2834        let selected = task.output_companions();
2835        if declared.len() != selected.len()
2836            || declared.iter().zip(selected).any(|(declared, selected)| {
2837                declared.name() != selected.name()
2838                    || declared.role() != selected.role()
2839                    || declared.logical_shape() != selected.logical_shape()
2840                    || !(declared.owner() == selected.owner()
2841                        || matches!(
2842                            (declared.owner(), selected.owner()),
2843                            (
2844                                ParameterGroupOwner::StaticAnyOf(declared_roles),
2845                                ParameterGroupOwner::StaticRole(selected_role)
2846                            ) if declared_roles.iter().any(|role| role == selected_role)
2847                        ))
2848            })
2849        {
2850            return Err(ReplicatedTextContractError::invalid(format!(
2851                "partition parameter companions for {:?} differ from authoritative selection: constructed={:?}, selected={:?}",
2852                task.name(), declared, selected,
2853            )));
2854        }
2855    }
2856    if !companions.is_empty() {
2857        return Err(ReplicatedTextContractError::invalid(format!(
2858            "architecture companions name missing primary tasks: {:?}",
2859            companions.keys().collect::<Vec<_>>()
2860        )));
2861    }
2862    let mut projected = Vec::new();
2863    for task in tasks {
2864        let topology_target = topology_targets
2865            .get(task.name())
2866            .expect("every task has one validated topology target");
2867        let emitted = std::iter::once(topology_target.as_str())
2868            .chain(
2869                task.output_companions()
2870                    .iter()
2871                    .map(ReplicatedTextOutputCompanion::name),
2872            )
2873            .collect::<Vec<_>>();
2874        let local = emitted
2875            .iter()
2876            .filter(|target| owned_targets.contains(**target))
2877            .count();
2878        match local {
2879            0 => {}
2880            count if count == emitted.len() => projected.push(task),
2881            count => {
2882                return Err(ReplicatedTextContractError::invalid(format!(
2883                    "materialization task {:?} would emit {count} of {} outputs into this partition",
2884                    task.name(),
2885                    emitted.len()
2886                )));
2887            }
2888        }
2889    }
2890    Ok(projected)
2891}
2892
2893/// Parameter visitation order is an implementation detail of a local module,
2894/// while an architecture parameter group is an atomic, target-keyed contract.
2895/// Compare that contract without making otherwise-identical local ownership
2896/// depend on whether a backend-neutral module visits a bias before its weight.
2897fn parameter_groups_have_same_members(
2898    left: &ParameterGroupSpec,
2899    right: &ParameterGroupSpec,
2900) -> bool {
2901    left.logical_name() == right.logical_name()
2902        && left.role() == right.role()
2903        && left.partition_units() == right.partition_units()
2904        && left.members().len() == right.members().len()
2905        && left.members().iter().all(|left_member| {
2906            right.members().iter().any(|right_member| {
2907                left_member.target() == right_member.target()
2908                    && left_member.global_shape() == right_member.global_shape()
2909                    && left_member.sharding() == right_member.sharding()
2910                    && left_member.linear_companion() == right_member.linear_companion()
2911                    && left_member.linear_companion_of() == right_member.linear_companion_of()
2912            })
2913        })
2914}
2915
2916impl SelectedParameterRealization {
2917    /// Returns the canonical logical identity.
2918    pub fn name(&self) -> &str {
2919        &self.name
2920    }
2921    /// Returns admitted physical source identities.
2922    pub fn sources(&self) -> &[String] {
2923        &self.sources
2924    }
2925    /// Returns the exact selected shard and multi-output provenance.
2926    pub fn physical_sources(&self) -> &[ReplicatedTextPhysicalSource] {
2927        &self.physical_sources
2928    }
2929    /// Returns the admitted source encoding.
2930    pub const fn source_encoding(&self) -> &SourceTensorEncoding {
2931        &self.source_encoding
2932    }
2933    /// Returns the selected executable format.
2934    pub const fn executable(&self) -> LinearFormat {
2935        self.executable
2936    }
2937    /// Returns the selected backend lowering kind.
2938    pub const fn lowering(&self) -> WeightLoweringKind {
2939        self.lowering
2940    }
2941}
2942
2943/// Selected physical realization of one exact semantic state component.
2944#[derive(Debug, Clone, Eq, PartialEq)]
2945pub struct SelectedStateComponentRealization {
2946    storage_dtype: StateStorageDtype,
2947    layer: usize,
2948    component: StateComponentPolicy,
2949    placement: StateComponentPlacement,
2950}
2951
2952impl SelectedStateComponentRealization {
2953    /// Exact native scalar representation admitted for this component.
2954    pub const fn storage_dtype(&self) -> StateStorageDtype {
2955        self.storage_dtype
2956    }
2957
2958    /// Returns the architecture-global state layer.
2959    pub const fn layer(&self) -> usize {
2960        self.layer
2961    }
2962
2963    /// Returns the exact architecture-declared component contract.
2964    pub const fn component(&self) -> &StateComponentPolicy {
2965        &self.component
2966    }
2967
2968    /// Returns the selected physical placement.
2969    pub const fn placement(&self) -> StateComponentPlacement {
2970        self.placement
2971    }
2972}
2973
2974/// Authoritative mutable-state realization selected before allocation.
2975#[derive(Debug, Clone, Eq, PartialEq)]
2976pub struct SelectedStateRealization {
2977    floating_dtype: Option<StateStorageDtype>,
2978    layout: StateLayout,
2979    access: ReplicatedTextStateAccess,
2980    policy: CacheResidencyPolicy,
2981    components: Vec<SelectedStateComponentRealization>,
2982    checkpoint: bool,
2983    rollback: bool,
2984    reset: bool,
2985    prompt_cache: bool,
2986    observation_retention: bool,
2987}
2988
2989impl SelectedStateRealization {
2990    /// Native representation selected from the architecture's floating-state source.
2991    pub const fn floating_dtype(&self) -> Option<StateStorageDtype> {
2992        self.floating_dtype
2993    }
2994
2995    /// Returns the exact architecture-owned state layout.
2996    pub const fn layout(&self) -> &StateLayout {
2997        &self.layout
2998    }
2999
3000    /// Selects the exact rank-local state interval while preserving global ownership proof.
3001    ///
3002    /// Component ordinals are rebased to the local layout consumed by a rank-local runtime;
3003    /// prompt-cache identity retains the global offset separately through [`crate::PartitionState`].
3004    pub fn for_partition(
3005        &self,
3006        partition: &crate::PartitionState,
3007    ) -> Result<Self, ReplicatedTextContractError> {
3008        let range = partition.global_layers();
3009        let expected = self
3010            .layout
3011            .slice(range.clone())
3012            .map_err(|error| ReplicatedTextContractError::invalid(error.to_string()))?;
3013        if &expected != partition.layout() {
3014            return Err(ReplicatedTextContractError::invalid(
3015                "partition state layout differs from the selected global interval",
3016            ));
3017        }
3018        let components = self
3019            .components
3020            .iter()
3021            .filter(|component| range.contains(&component.layer))
3022            .cloned()
3023            .map(|mut component| {
3024                component.layer -= range.start;
3025                component
3026            })
3027            .collect::<Vec<_>>();
3028        let expected_components = (0..partition.layout().len())
3029            .map(|layer| {
3030                partition
3031                    .layout()
3032                    .components(layer)
3033                    .expect("validated local state layout contains every layer")
3034                    .len()
3035            })
3036            .sum::<usize>();
3037        if components.len() != expected_components {
3038            return Err(ReplicatedTextContractError::invalid(
3039                "partition state components differ from the selected global interval",
3040            ));
3041        }
3042        Ok(Self {
3043            floating_dtype: self.floating_dtype,
3044            layout: partition.layout().clone(),
3045            access: self.access,
3046            policy: self.policy.clone(),
3047            components,
3048            checkpoint: self.checkpoint,
3049            rollback: self.rollback,
3050            reset: self.reset,
3051            prompt_cache: self.prompt_cache,
3052            observation_retention: self.observation_retention,
3053        })
3054    }
3055
3056    /// Selects a rank-local interval whose tensor-parallel component shapes were authored by the
3057    /// validated architecture partition.
3058    ///
3059    /// Pipeline ownership must still name the same global layer interval. Tensor-parallel
3060    /// geometry may narrow fixed dimensions, but it cannot change component roles, dtype,
3061    /// residency, presence, ordering, or selected physical placement.
3062    pub fn for_partitioned_geometry(
3063        &self,
3064        partition: &crate::PartitionState,
3065    ) -> Result<Self, ReplicatedTextContractError> {
3066        let range = partition.global_layers();
3067        let global = self
3068            .layout
3069            .slice(range.clone())
3070            .map_err(|error| ReplicatedTextContractError::invalid(error.to_string()))?;
3071        if global.len() != partition.layout().len() {
3072            return Err(ReplicatedTextContractError::invalid(
3073                "partition state layer count differs from the selected global interval",
3074            ));
3075        }
3076        let mut components = Vec::new();
3077        for local_layer in 0..partition.layout().len() {
3078            let global_components = global
3079                .components(local_layer)
3080                .expect("validated selected state contains every local layer");
3081            let local_components = partition
3082                .layout()
3083                .components(local_layer)
3084                .expect("validated partition state contains every local layer");
3085            if global_components.len() != local_components.len() {
3086                return Err(ReplicatedTextContractError::invalid(
3087                    "partition state component count differs from selected state",
3088                ));
3089            }
3090            let global_layer = range.start + local_layer;
3091            let selected_components = self
3092                .components
3093                .iter()
3094                .filter(|component| component.layer == global_layer)
3095                .collect::<Vec<_>>();
3096            if selected_components.len() != local_components.len() {
3097                return Err(ReplicatedTextContractError::invalid(
3098                    "partition state components differ from the selected global interval",
3099                ));
3100            }
3101            for ((global_policy, local_policy), selected) in global_components
3102                .iter()
3103                .zip(local_components)
3104                .zip(selected_components)
3105            {
3106                if global_policy.role() != local_policy.role()
3107                    || global_policy.dtype() != local_policy.dtype()
3108                    || global_policy.residency() != local_policy.residency()
3109                    || global_policy.presence() != local_policy.presence()
3110                    || selected.component != *global_policy
3111                {
3112                    return Err(ReplicatedTextContractError::invalid(
3113                        "partition state component semantics differ from selected state",
3114                    ));
3115                }
3116                components.push(SelectedStateComponentRealization {
3117                    layer: local_layer,
3118                    component: local_policy.clone(),
3119                    storage_dtype: selected.storage_dtype,
3120                    placement: selected.placement,
3121                });
3122            }
3123        }
3124        Ok(Self {
3125            floating_dtype: self.floating_dtype,
3126            layout: partition.layout().clone(),
3127            access: self.access,
3128            policy: self.policy.clone(),
3129            components,
3130            checkpoint: self.checkpoint,
3131            rollback: self.rollback,
3132            reset: self.reset,
3133            prompt_cache: self.prompt_cache,
3134            observation_retention: self.observation_retention,
3135        })
3136    }
3137
3138    /// Returns the state-access semantics selected for typed traversal.
3139    pub const fn access(&self) -> ReplicatedTextStateAccess {
3140        self.access
3141    }
3142
3143    /// Returns the selected residency policy.
3144    pub const fn policy(&self) -> &CacheResidencyPolicy {
3145        &self.policy
3146    }
3147
3148    /// Returns exact selected component realizations in layer/component order.
3149    pub fn components(&self) -> &[SelectedStateComponentRealization] {
3150        &self.components
3151    }
3152
3153    /// Returns whether state checkpoints are selected.
3154    pub const fn checkpoint(&self) -> bool {
3155        self.checkpoint
3156    }
3157
3158    /// Returns whether checkpoint rollback is selected.
3159    pub const fn rollback(&self) -> bool {
3160        self.rollback
3161    }
3162
3163    /// Returns whether complete reset is selected.
3164    pub const fn reset(&self) -> bool {
3165        self.reset
3166    }
3167
3168    /// Returns whether prompt-cache persistence is selected.
3169    pub const fn prompt_cache(&self) -> bool {
3170        self.prompt_cache
3171    }
3172
3173    /// Returns whether observation retains every live component.
3174    pub const fn observation_retention(&self) -> bool {
3175        self.observation_retention
3176    }
3177}
3178
3179/// Authoritative realization selected before architecture or payload construction.
3180#[derive(Debug, Clone, Eq, PartialEq)]
3181pub struct SelectedReplicatedTextRealization {
3182    max_cached_shards: usize,
3183    requirements: ReplicatedTextRequirements,
3184    /// Exact selected execution topology.
3185    topology: ParallelTopology,
3186    /// Selected ordinary parameter residency.
3187    residency: LayerWeightResidency,
3188    /// Selected exact mutable-state implementation.
3189    state: SelectedStateRealization,
3190    /// Exact per-parameter source, executable format, and lowering.
3191    parameters: Vec<SelectedParameterRealization>,
3192    /// Exact task/companion topology selected before stores or native modules exist.
3193    materialization_tasks: Vec<ReplicatedTextMaterializationTask>,
3194    /// Exact additive auxiliary parameter selections.
3195    auxiliary_parameters: Vec<SelectedParameterRealization>,
3196    /// Exact additive auxiliary task/companion topology.
3197    auxiliary_materialization_tasks: Vec<ReplicatedTextMaterializationTask>,
3198    /// Required observation facilities admitted by the backend.
3199    session: SessionCapabilities,
3200    /// Prompt-cache persistence is selected for this lifecycle.
3201    prompt_cache: bool,
3202    /// Exact completion ownership selected for this lifecycle.
3203    exact_completion: bool,
3204    grouped_operations: Vec<GroupedOperationRequirement>,
3205}
3206
3207impl SelectedReplicatedTextRealization {
3208    /// Returns the source reader-cache limit retained through selection.
3209    pub const fn max_cached_shards(&self) -> usize {
3210        self.max_cached_shards
3211    }
3212    /// Returns the exact architecture/artifact requirements selected together.
3213    pub const fn requirements(&self) -> &ReplicatedTextRequirements {
3214        &self.requirements
3215    }
3216    /// Returns the exact selected topology.
3217    pub const fn topology(&self) -> ParallelTopology {
3218        self.topology
3219    }
3220    /// Returns selected weight residency.
3221    pub const fn residency(&self) -> LayerWeightResidency {
3222        self.residency
3223    }
3224    /// Returns the authoritative selected mutable-state realization.
3225    pub const fn state(&self) -> &SelectedStateRealization {
3226        &self.state
3227    }
3228    /// Returns exact per-parameter realizations.
3229    pub fn parameters(&self) -> &[SelectedParameterRealization] {
3230        &self.parameters
3231    }
3232    /// Returns the authoritative exact materialization task sequence.
3233    pub fn materialization_tasks(&self) -> &[ReplicatedTextMaterializationTask] {
3234        &self.materialization_tasks
3235    }
3236    /// Returns exact additive auxiliary parameter selections.
3237    pub fn auxiliary_parameters(&self) -> &[SelectedParameterRealization] {
3238        &self.auxiliary_parameters
3239    }
3240    /// Returns exact additive auxiliary task/companion topology.
3241    pub fn auxiliary_materialization_tasks(&self) -> &[ReplicatedTextMaterializationTask] {
3242        &self.auxiliary_materialization_tasks
3243    }
3244    /// Returns selected session facilities.
3245    pub const fn session(&self) -> SessionCapabilities {
3246        self.session
3247    }
3248    /// Returns whether prompt-cache persistence was selected.
3249    pub const fn prompt_cache(&self) -> bool {
3250        self.prompt_cache
3251    }
3252    /// Returns whether exact completion ownership was selected.
3253    pub const fn exact_completion(&self) -> bool {
3254        self.exact_completion
3255    }
3256    /// Returns selected grouped operation mechanisms.
3257    pub fn grouped_operations(&self) -> &[GroupedOperationRequirement] {
3258        &self.grouped_operations
3259    }
3260}
3261
3262/// Complete fail-closed selection diagnostic.
3263#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
3264#[error("replicated text realization is unsupported: {issues}", issues = .issues.join("; "))]
3265pub struct ReplicatedTextSelectionError {
3266    issues: Vec<String>,
3267}
3268
3269impl ReplicatedTextSelectionError {
3270    /// Every missing semantic or mechanism requirement in stable order.
3271    pub fn issues(&self) -> &[String] {
3272        &self.issues
3273    }
3274}
3275
3276/// Deterministically selects one realization without constructing backend payloads.
3277pub fn select_replicated_text_realization(
3278    requirements: &ReplicatedTextRequirements,
3279    request: &ReplicatedTextSelectionRequest,
3280    capabilities: &BackendMechanismCapabilities,
3281) -> Result<SelectedReplicatedTextRealization, ReplicatedTextSelectionError> {
3282    let mut issues = Vec::new();
3283    if request
3284        .topology
3285        .is_some_and(|topology| !topology.is_replicated())
3286    {
3287        issues.push("replicated execution topology".into());
3288    }
3289    if !capabilities.operators.contains(requirements.operators) {
3290        issues.extend(
3291            capabilities
3292                .operators
3293                .missing_capability_names(requirements.operators)
3294                .into_iter()
3295                .map(|name| format!("neural operation {name}")),
3296        );
3297    }
3298    for operation in &requirements.grouped_operations {
3299        if !capabilities.grouped_operations.contains(operation) {
3300            issues.push(format!("grouped operation {operation:?}"));
3301        }
3302    }
3303    let residency_mechanism = match request.residency {
3304        LayerWeightResidency::FullyResident => WeightResidencyMechanism::Resident,
3305        LayerWeightResidency::LayerwiseHost(_) => WeightResidencyMechanism::Windowed,
3306        LayerWeightResidency::DenseDiskStream(_) => WeightResidencyMechanism::DiskStreamed,
3307    };
3308    if !capabilities
3309        .weight_residencies
3310        .contains(&residency_mechanism)
3311    {
3312        issues.push(format!("weight residency {residency_mechanism:?}"));
3313    }
3314    let floating_dtype = match capabilities.state.floating_state_dtype() {
3315        Some((source, dtype))
3316            if Some(source) == requirements.floating_state_source() && dtype.is_floating() =>
3317        {
3318            Some(dtype)
3319        }
3320        Some(_) => {
3321            issues.push("floating-state dtype support differs from the selected source".into());
3322            None
3323        }
3324        None => None,
3325    };
3326    let mut state_components = Vec::new();
3327    for layer in 0..requirements.state_layout.len() {
3328        for component in requirements
3329            .state_layout
3330            .components(layer)
3331            .expect("state layout exposes every validated layer")
3332        {
3333            let Some(storage_dtype) = StateStorageDtype::resolve(component.dtype(), floating_dtype)
3334            else {
3335                issues.push(format!(
3336                    "state component {} at layer {layer} has no selected floating storage dtype",
3337                    component.role().stable_name()
3338                ));
3339                continue;
3340            };
3341            let matches = capabilities
3342                .state
3343                .components
3344                .iter()
3345                .filter(|mechanism| mechanism.layer == layer && mechanism.component == *component)
3346                .collect::<Vec<_>>();
3347            let role = component.role().stable_name();
3348            match matches.as_slice() {
3349                [mechanism] => match mechanism.placement(&request.state) {
3350                    Some(placement) if placement_is_compatible(component, &request.state, placement) => {
3351                        state_components.push(SelectedStateComponentRealization {
3352                            layer,
3353                            component: component.clone(),
3354                            storage_dtype,
3355                            placement,
3356                        });
3357                    }
3358                    Some(placement) => issues.push(format!(
3359                        "state component {role} at layer {layer} has incompatible {placement:?} placement for {:?} and {:?} residency",
3360                        request.state,
3361                        component.residency()
3362                    )),
3363                    None => issues.push(format!(
3364                        "state component {role} at layer {layer} for {:?}",
3365                        request.state
3366                    )),
3367                },
3368                [] => issues.push(format!(
3369                    "state component {role} at layer {layer} with shape {:?} and dtype {:?}",
3370                    component.shape(),
3371                    component.dtype()
3372                )),
3373                _ => issues.push(format!(
3374                    "unique state component mechanism {role} at layer {layer}"
3375                )),
3376            }
3377        }
3378    }
3379    for (supported, name) in [
3380        (capabilities.state.checkpoint, "state checkpoint"),
3381        (capabilities.state.rollback, "state rollback"),
3382        (capabilities.state.reset, "state reset"),
3383    ] {
3384        if !supported {
3385            issues.push(name.into());
3386        }
3387    }
3388    if request.prompt_cache && !capabilities.state.prompt_cache {
3389        issues.push("state prompt-cache persistence".into());
3390    }
3391    if (request.session.output_observation() || request.session.activation_inspection())
3392        && !capabilities.state.observation_retention
3393    {
3394        issues.push("state observation retention".into());
3395    }
3396    for (required, supported, name) in [
3397        (
3398            request.session.persistent_cache(),
3399            capabilities.session.persistent_cache(),
3400            "persistent_cache",
3401        ),
3402        (
3403            request.session.output_observation(),
3404            capabilities.session.output_observation(),
3405            "output_observation",
3406        ),
3407        (
3408            request.session.activation_inspection(),
3409            capabilities.session.activation_inspection(),
3410            "activation_inspection",
3411        ),
3412    ] {
3413        if required && !supported {
3414            issues.push(format!("session capability {name}"));
3415        }
3416    }
3417    if request.prompt_cache && !capabilities.prompt_cache {
3418        issues.push("prompt-cache persistence".into());
3419    }
3420    if request.exact_completion && !capabilities.exact_completion {
3421        issues.push("exact completion ownership".into());
3422    }
3423
3424    let mut parameters = Vec::with_capacity(requirements.parameters.len());
3425    let mut auxiliary_parameters = Vec::with_capacity(requirements.auxiliary_parameters.len());
3426    let mut names = BTreeSet::new();
3427    for (parameter, auxiliary) in requirements
3428        .parameters
3429        .iter()
3430        .map(|parameter| (parameter, false))
3431        .chain(
3432            requirements
3433                .auxiliary_parameters
3434                .iter()
3435                .map(|parameter| (parameter, true)),
3436        )
3437    {
3438        if parameter.name.trim().is_empty() || !names.insert(parameter.name.as_str()) {
3439            issues.push(format!(
3440                "unique nonempty logical parameter identity {:?}",
3441                parameter.name
3442            ));
3443            continue;
3444        }
3445        if !parameter.has_lowering_source() {
3446            continue;
3447        }
3448        let native_candidate = || {
3449            parameter
3450                .lowering_descriptor(parameter.native_executable)
3451                .map(|descriptor| (parameter.native_executable, descriptor))
3452        };
3453        let candidate = match request.quantization {
3454            Some(request) => parameter
3455                .transform_target(request)
3456                .and_then(|target| match target {
3457                    Some(target) => Ok((target.executable(), target.descriptor().clone())),
3458                    None => native_candidate(),
3459                }),
3460            None => native_candidate(),
3461        };
3462        let (executable, descriptor) = match candidate {
3463            Ok(candidate) => candidate,
3464            Err(error) => {
3465                issues.push(error.to_string());
3466                issues.push(format!(
3467                    "architecture transform {:?} for {:?}",
3468                    request.quantization, parameter.name
3469                ));
3470                continue;
3471            }
3472        };
3473        let Some(lowering) = capabilities
3474            .weight_lowerings
3475            .iter()
3476            .find(|lowering| lowering.descriptor == descriptor)
3477        else {
3478            issues.push(format!(
3479                "weight lowering {:?} -> {:?} for {:?} with descriptor {:?}",
3480                parameter.source_encoding, executable, parameter.name, descriptor
3481            ));
3482            continue;
3483        };
3484        let selected_parameter = SelectedParameterRealization {
3485            name: parameter.name.clone(),
3486            sources: parameter.sources.clone(),
3487            physical_sources: parameter.physical_sources.clone(),
3488            source_encoding: parameter
3489                .source_encoding
3490                .clone()
3491                .expect("physical parameter has a source encoding"),
3492            executable,
3493            lowering: match (&parameter.presence, lowering.kind) {
3494                (
3495                    ReplicatedTextParameterPresence::Derived { .. },
3496                    WeightLoweringKind::Transform | WeightLoweringKind::DerivedTransform,
3497                ) => WeightLoweringKind::DerivedTransform,
3498                (ReplicatedTextParameterPresence::Derived { .. }, _) => WeightLoweringKind::Derived,
3499                (_, kind) => kind,
3500            },
3501        };
3502        if auxiliary {
3503            auxiliary_parameters.push(selected_parameter);
3504        } else {
3505            parameters.push(selected_parameter);
3506        }
3507    }
3508    if !issues.is_empty() {
3509        return Err(ReplicatedTextSelectionError { issues });
3510    }
3511    let mut selected = SelectedReplicatedTextRealization {
3512        max_cached_shards: request.max_cached_shards,
3513        requirements: requirements.clone(),
3514        topology: request
3515            .topology
3516            .unwrap_or_else(|| ParallelTopology::new(1, 1, 1, 1).expect("replicated topology")),
3517        residency: request.residency,
3518        state: SelectedStateRealization {
3519            floating_dtype,
3520            layout: requirements.state_layout.clone(),
3521            access: requirements.state_access,
3522            policy: request.state.clone(),
3523            components: state_components,
3524            checkpoint: true,
3525            rollback: true,
3526            reset: true,
3527            prompt_cache: request.prompt_cache,
3528            observation_retention: request.session.output_observation()
3529                || request.session.activation_inspection(),
3530        },
3531        parameters,
3532        materialization_tasks: Vec::new(),
3533        auxiliary_parameters,
3534        auxiliary_materialization_tasks: Vec::new(),
3535        session: request.session,
3536        prompt_cache: request.prompt_cache,
3537        exact_completion: request.exact_completion,
3538        grouped_operations: requirements.grouped_operations.clone(),
3539    };
3540    selected.materialization_tasks = build_replicated_text_materialization_tasks(&selected)
3541        .map_err(|error| ReplicatedTextSelectionError {
3542            issues: vec![error.to_string()],
3543        })?;
3544    selected.auxiliary_materialization_tasks = build_materialization_tasks(
3545        selected.requirements(),
3546        selected.requirements().auxiliary_parameters(),
3547        &selected.auxiliary_parameters,
3548    )
3549    .map_err(|error| ReplicatedTextSelectionError {
3550        issues: vec![error.to_string()],
3551    })?;
3552    Ok(selected)
3553}
3554
3555pub(crate) fn placement_is_compatible(
3556    component: &StateComponentPolicy,
3557    policy: &CacheResidencyPolicy,
3558    placement: StateComponentPlacement,
3559) -> bool {
3560    use eredu_core::cache::StateResidencyClass;
3561
3562    let expected = match (policy, component.residency()) {
3563        (CacheResidencyPolicy::Device, _) => StateComponentPlacement::Device,
3564        (CacheResidencyPolicy::Paged(_), StateResidencyClass::SealablePaged) => {
3565            StateComponentPlacement::Paged
3566        }
3567        (
3568            CacheResidencyPolicy::Paged(_),
3569            StateResidencyClass::AlwaysDeviceMutable | StateResidencyClass::LayerScopedOffloadable,
3570        ) => StateComponentPlacement::Device,
3571    };
3572    placement == expected
3573}
3574
3575#[cfg(test)]
3576mod tests {
3577    use super::*;
3578    use crate::{
3579        ArchitectureGroupKind, ArchitectureGroupPlacement, ArchitectureGroupTransport,
3580        ArchitectureMergeDestination, ArchitectureParameterDescription, ArchitecturePartition,
3581        ArchitectureStatePartitionPlan, ArchitectureStatePartitionRule, DenseDiskStreamLoadOptions,
3582        ExecutionGroupSpec, ExecutionUnitLayout, LayerwiseLoadOptions, MemberSharding,
3583        NoAuxiliaryBoundarySchema, OwnedParameterGroupSpec, ParameterGroupSpec,
3584        ParameterMemberSpec, ParameterRole, PartitionOwnership, StateLayout,
3585    };
3586    use eredu_checkpoint::{AffineQuantization, StoredDtype};
3587    use eredu_core::{
3588        cache::{
3589            LayerCachePolicy, MutableStateResidency, StateTensorDimension, StateTensorDtype,
3590            StateTensorPolicy, StateTensorRole,
3591        },
3592        AttentionPolicy, LayerSchedule,
3593    };
3594
3595    fn paged_state() -> CacheResidencyPolicy {
3596        CacheResidencyPolicy::Paged(
3597            crate::PagedCacheOptions::new(4, 1 << 20, 1 << 20, 1)
3598                .unwrap()
3599                .with_full_attention(true),
3600        )
3601    }
3602
3603    fn physical_source(name: &str) -> ReplicatedTextPhysicalSource {
3604        ReplicatedTextPhysicalSource::new(
3605            name,
3606            name,
3607            "/checkpoint/model.safetensors",
3608            name,
3609            SourceTensorEncoding::Safetensors(StoredDtype::F16),
3610            2,
3611        )
3612        .unwrap()
3613    }
3614
3615    fn requirements() -> ReplicatedTextRequirements {
3616        let graph =
3617            ExecutionGraph::new(vec![ExecutionGroupSpec::root("decoder")], "decoder").unwrap();
3618        let execution_units = ExecutionUnitLayout::new(&graph, [1]).unwrap();
3619        ReplicatedTextRequirements::new(
3620            "test.replicated-text",
3621            NeuralOperatorCapabilities::EXP,
3622            graph,
3623            execution_units,
3624            vec![ArchitectureGroupTransport {
3625                placement: ArchitectureGroupPlacement::Pipeline,
3626                kind: ArchitectureGroupKind::Decoder,
3627                first_owner_static_roles: vec!["embedding".into()],
3628                last_owner_static_roles: vec!["output".into()],
3629                merge_destination: ArchitectureMergeDestination::LastOwner,
3630                parallel_subgroup: None,
3631                request_optional: false,
3632            }],
3633            StateLayout::new(
3634                LayerSchedule::new(
3635                    1,
3636                    vec![LayerCachePolicy::key_value(AttentionPolicy::Full, 1, 8).unwrap()],
3637                )
3638                .unwrap(),
3639            )
3640            .unwrap(),
3641            ReplicatedTextStateAccess::KeyValue,
3642            vec![
3643                ReplicatedTextParameterRequirement::new(
3644                    "model.layers.0.mlp.weight",
3645                    vec!["blk.0.ffn.weight".into()],
3646                    vec![physical_source("blk.0.ffn.weight")],
3647                    Vec::new(),
3648                    Some(SourceTensorEncoding::Safetensors(StoredDtype::F16)),
3649                    Some(vec![64, 64]),
3650                    vec![64, 64],
3651                    LinearFormat::Dense,
3652                    ReplicatedTextParameterRole::LinearWeight,
3653                    ReplicatedTextParameterOwner::ExecutionUnit {
3654                        group: "decoder".into(),
3655                        unit: 0,
3656                    },
3657                    ReplicatedTextParameterPresence::Required,
3658                    ParameterTransformConstraint::Linear { packed_axis: 1 },
3659                )
3660                .and_then(|requirement| {
3661                    requirement.with_transform_companions(
3662                        "model.layers.0.mlp.scales",
3663                        "model.layers.0.mlp.biases",
3664                    )
3665                })
3666                .unwrap(),
3667                ReplicatedTextParameterRequirement::new(
3668                    "model.layers.0.mlp.bias",
3669                    Vec::new(),
3670                    Vec::new(),
3671                    Vec::new(),
3672                    None,
3673                    None,
3674                    vec![64],
3675                    LinearFormat::Dense,
3676                    ReplicatedTextParameterRole::LinearBias,
3677                    ReplicatedTextParameterOwner::ExecutionUnit {
3678                        group: "decoder".into(),
3679                        unit: 0,
3680                    },
3681                    ReplicatedTextParameterPresence::OptionalAbsent,
3682                    ParameterTransformConstraint::None,
3683                )
3684                .unwrap(),
3685                ReplicatedTextParameterRequirement::new(
3686                    "model.layers.0.norm.weight",
3687                    vec!["blk.0.norm.weight".into()],
3688                    vec![physical_source("blk.0.norm.weight")],
3689                    Vec::new(),
3690                    Some(SourceTensorEncoding::Safetensors(StoredDtype::F16)),
3691                    Some(vec![64]),
3692                    vec![64],
3693                    LinearFormat::Dense,
3694                    ReplicatedTextParameterRole::Normalization,
3695                    ReplicatedTextParameterOwner::ExecutionUnit {
3696                        group: "decoder".into(),
3697                        unit: 0,
3698                    },
3699                    ReplicatedTextParameterPresence::Required,
3700                    ParameterTransformConstraint::None,
3701                )
3702                .unwrap(),
3703            ],
3704        )
3705        .unwrap()
3706        .with_floating_state_source(TensorDtype::F16)
3707    }
3708
3709    #[test]
3710    fn requirements_reject_unit_layout_from_an_equally_sized_different_graph() {
3711        let baseline = requirements();
3712        let other_graph =
3713            ExecutionGraph::new(vec![ExecutionGroupSpec::root("mutated")], "mutated").unwrap();
3714        let other_layout = ExecutionUnitLayout::new(&other_graph, [1]).unwrap();
3715        let error = ReplicatedTextRequirements::new(
3716            baseline.architecture_identity.clone(),
3717            baseline.operators,
3718            baseline.execution_graph.clone(),
3719            other_layout,
3720            baseline.group_transports.clone(),
3721            baseline.state_layout.clone(),
3722            baseline.state_access,
3723            baseline.parameters.clone(),
3724        )
3725        .unwrap_err();
3726        assert!(error.to_string().contains("layout group identities differ"));
3727    }
3728
3729    #[test]
3730    fn parameter_requirement_preserves_every_admitted_alias() {
3731        let requirement = ReplicatedTextParameterRequirement::new(
3732            "model.layers.0.mlp.weight",
3733            vec!["released.layers.0.mlp.weight".into()],
3734            vec![physical_source("released.layers.0.mlp.weight")],
3735            vec![
3736                "legacy.layers.0.mlp.weight".into(),
3737                "vendor.layers.0.mlp.weight".into(),
3738            ],
3739            Some(SourceTensorEncoding::Safetensors(StoredDtype::F16)),
3740            Some(vec![64, 64]),
3741            vec![64, 64],
3742            LinearFormat::Dense,
3743            ReplicatedTextParameterRole::LinearWeight,
3744            ReplicatedTextParameterOwner::ExecutionUnit {
3745                group: "decoder".into(),
3746                unit: 0,
3747            },
3748            ReplicatedTextParameterPresence::Required,
3749            ParameterTransformConstraint::Linear { packed_axis: 1 },
3750        )
3751        .unwrap();
3752
3753        assert_eq!(
3754            requirement.aliases(),
3755            ["legacy.layers.0.mlp.weight", "vendor.layers.0.mlp.weight"]
3756        );
3757        assert_eq!(requirement.sources(), ["released.layers.0.mlp.weight"]);
3758
3759        let absent_bias = ReplicatedTextParameterRequirement::new(
3760            "model.layers.0.mlp.bias",
3761            Vec::new(),
3762            Vec::new(),
3763            vec!["released.layers.0.mlp.bias".into()],
3764            None,
3765            None,
3766            vec![64],
3767            LinearFormat::Dense,
3768            ReplicatedTextParameterRole::LinearBias,
3769            ReplicatedTextParameterOwner::ExecutionUnit {
3770                group: "decoder".into(),
3771                unit: 0,
3772            },
3773            ReplicatedTextParameterPresence::OptionalAbsent,
3774            ParameterTransformConstraint::None,
3775        )
3776        .unwrap();
3777        assert_eq!(
3778            absent_bias.presence(),
3779            &ReplicatedTextParameterPresence::OptionalAbsent
3780        );
3781        assert!(absent_bias.sources().is_empty());
3782        assert_eq!(
3783            absent_bias.transform_constraint(),
3784            ParameterTransformConstraint::None
3785        );
3786    }
3787
3788    #[test]
3789    fn scalar_parameter_requirement_preserves_rank_zero_geometry() {
3790        let requirement = ReplicatedTextParameterRequirement::new(
3791            "model.audio_tower.input_max",
3792            vec!["model.audio_tower.input_max".into()],
3793            vec![physical_source("model.audio_tower.input_max")],
3794            Vec::new(),
3795            Some(SourceTensorEncoding::Safetensors(StoredDtype::F32)),
3796            Some(Vec::new()),
3797            Vec::new(),
3798            LinearFormat::Dense,
3799            ReplicatedTextParameterRole::Other,
3800            ReplicatedTextParameterOwner::StaticRole("audio".into()),
3801            ReplicatedTextParameterPresence::Required,
3802            ParameterTransformConstraint::None,
3803        )
3804        .unwrap();
3805
3806        let descriptor = requirement
3807            .lowering_descriptor(LinearFormat::Dense)
3808            .unwrap();
3809        assert!(descriptor.physical_shape().is_empty());
3810        assert!(descriptor.logical_shape().is_empty());
3811        assert_eq!(descriptor.packed_axis(), None);
3812    }
3813
3814    #[test]
3815    fn physical_provenance_distinguishes_outputs_from_one_sharded_tensor() {
3816        let shard = "/checkpoint/model-00002-of-00003.gguf";
3817        let weight = ReplicatedTextPhysicalSource::new(
3818            "model.layers.0.gate_proj.weight",
3819            "blk.0.ffn_gate.weight",
3820            shard,
3821            "blk.0.ffn_gate.weight",
3822            SourceTensorEncoding::Safetensors(StoredDtype::F16),
3823            2,
3824        )
3825        .unwrap();
3826        let scales = ReplicatedTextPhysicalSource::new(
3827            "model.layers.0.gate_proj.scales",
3828            "blk.0.ffn_gate.weight",
3829            shard,
3830            "blk.0.ffn_gate.scales",
3831            SourceTensorEncoding::Safetensors(StoredDtype::F16),
3832            2,
3833        )
3834        .unwrap();
3835        assert_eq!(weight.tensor(), scales.tensor());
3836        assert_eq!(weight.shard(), scales.shard());
3837        assert_ne!(weight.output(), scales.output());
3838    }
3839
3840    fn capabilities() -> BackendMechanismCapabilities {
3841        let source = SourceTensorEncoding::Safetensors(StoredDtype::F16);
3842        let requirements = requirements();
3843        let state = StateMechanismCapabilities::new(
3844            (0..requirements.state_layout().len()).flat_map(|layer| {
3845                requirements
3846                    .state_layout()
3847                    .components(layer)
3848                    .unwrap()
3849                    .iter()
3850                    .cloned()
3851                    .map(move |component| {
3852                        let paged = match component.role() {
3853                            eredu_core::cache::StateComponentRole::AttentionKeys
3854                            | eredu_core::cache::StateComponentRole::AttentionValues
3855                            | eredu_core::cache::StateComponentRole::CompressedLatent
3856                            | eredu_core::cache::StateComponentRole::RotaryKeys => {
3857                                StateComponentPlacement::Paged
3858                            }
3859                            eredu_core::cache::StateComponentRole::Fixed(_) => {
3860                                StateComponentPlacement::Device
3861                            }
3862                        };
3863                        StateComponentMechanism::new(
3864                            layer,
3865                            component,
3866                            Some(StateComponentPlacement::Device),
3867                            Some(paged),
3868                        )
3869                    })
3870            }),
3871        )
3872        .with_floating_state_dtype(TensorDtype::F16, StateStorageDtype::F16)
3873        .with_transactions(true, true)
3874        .with_reset(true)
3875        .with_prompt_cache(true)
3876        .with_observation_retention(true);
3877        BackendMechanismCapabilities::new(
3878            NeuralOperatorCapabilities::EXP,
3879            vec![
3880                WeightLoweringCapability::new(
3881                    WeightLoweringDescriptor::new(
3882                        source.clone(),
3883                        LinearFormat::Dense,
3884                        vec![64, 64],
3885                        vec![64, 64],
3886                        Some(1),
3887                    )
3888                    .unwrap(),
3889                    WeightLoweringKind::Direct,
3890                ),
3891                WeightLoweringCapability::new(
3892                    WeightLoweringDescriptor::new(
3893                        source,
3894                        LinearFormat::Affine(AffineQuantization::new(64, 4).unwrap()),
3895                        vec![64, 64],
3896                        vec![64, 64],
3897                        Some(1),
3898                    )
3899                    .unwrap(),
3900                    WeightLoweringKind::Transform,
3901                ),
3902                WeightLoweringCapability::new(
3903                    WeightLoweringDescriptor::new(
3904                        SourceTensorEncoding::Safetensors(StoredDtype::F16),
3905                        LinearFormat::Dense,
3906                        vec![64],
3907                        vec![64],
3908                        None,
3909                    )
3910                    .unwrap(),
3911                    WeightLoweringKind::Direct,
3912                ),
3913            ],
3914            vec![
3915                WeightResidencyMechanism::Resident,
3916                WeightResidencyMechanism::Windowed,
3917                WeightResidencyMechanism::DiskStreamed,
3918            ],
3919            state,
3920        )
3921        .with_session(SessionCapabilities::new(true, true, true))
3922        .with_prompt_cache(true)
3923        .with_exact_completion(true)
3924    }
3925
3926    fn request(residency: LayerWeightResidency) -> ReplicatedTextSelectionRequest {
3927        ReplicatedTextSelectionRequest::new(residency, paged_state())
3928            .with_session(SessionCapabilities::new(true, true, true))
3929            .with_prompt_cache(true)
3930            .with_exact_completion(true)
3931    }
3932
3933    #[test]
3934    fn complete_requirements_are_invariant_across_all_caller_policy_dimensions() {
3935        let baseline = requirements();
3936        let disk = DenseDiskStreamLoadOptions::new(4096, 8192, 2, 1).unwrap();
3937        let requests = [
3938            ReplicatedTextSelectionRequest::new(
3939                LayerWeightResidency::FullyResident,
3940                CacheResidencyPolicy::Device,
3941            ),
3942            ReplicatedTextSelectionRequest::new(
3943                LayerWeightResidency::LayerwiseHost(LayerwiseLoadOptions::default()),
3944                paged_state(),
3945            )
3946            .with_topology(ParallelTopology::new(2, 1, 1, 1).unwrap())
3947            .with_quantization(QuantizationRequest::Affine {
3948                group_size: 64,
3949                bits: 4,
3950            })
3951            .with_session(SessionCapabilities::new(true, true, true))
3952            .with_prompt_cache(true)
3953            .with_exact_completion(true),
3954            ReplicatedTextSelectionRequest::new(
3955                LayerWeightResidency::DenseDiskStream(disk),
3956                CacheResidencyPolicy::Device,
3957            )
3958            .with_quantization(QuantizationRequest::MxFp4),
3959        ];
3960
3961        for _request in &requests {
3962            assert_eq!(requirements(), baseline);
3963        }
3964        assert_eq!(requests[0].state(), &CacheResidencyPolicy::Device);
3965        assert!(matches!(
3966            requests[1].residency(),
3967            LayerWeightResidency::LayerwiseHost(_)
3968        ));
3969        assert_eq!(requests[1].topology().unwrap().tensor(), 2);
3970        assert_eq!(
3971            requests[1].quantization(),
3972            Some(QuantizationRequest::Affine {
3973                group_size: 64,
3974                bits: 4,
3975            })
3976        );
3977        assert!(requests[1].prompt_cache());
3978        assert!(requests[1].exact_completion());
3979        assert!(requests[1].session().activation_inspection());
3980        assert_eq!(
3981            requests[2].residency(),
3982            LayerWeightResidency::DenseDiskStream(disk)
3983        );
3984        assert_eq!(requests[2].quantization(), Some(QuantizationRequest::MxFp4));
3985    }
3986
3987    #[test]
3988    fn partitioned_tasks_keep_encoded_companions_atomic_and_reject_split_groups() {
3989        let request = request(LayerWeightResidency::FullyResident).with_quantization(
3990            QuantizationRequest::Affine {
3991                group_size: 64,
3992                bits: 4,
3993            },
3994        );
3995        let selected =
3996            select_replicated_text_realization(&requirements(), &request, &capabilities()).unwrap();
3997        let graph = selected.requirements().execution_graph().clone();
3998        let layout = selected.requirements().execution_units().clone();
3999        let format = eredu_nn::LinearFormatSpec::affine(
4000            LinearFormat::Affine(AffineQuantization::new(64, 4).unwrap()),
4001            eredu_nn::ParameterSpec::trainable("model.layers.0.mlp.scales").unwrap(),
4002            eredu_nn::ParameterSpec::trainable("model.layers.0.mlp.biases").unwrap(),
4003        )
4004        .unwrap();
4005        let [physical] = crate::expand_linear_format_parameter_groups(
4006            vec![ParameterGroupSpec::new(
4007                "mlp",
4008                ParameterRole::FeedForwardIntermediate,
4009                [ParameterMemberSpec::new(
4010                    "model.layers.0.mlp.weight",
4011                    vec![64, 64],
4012                    MemberSharding::Replicated,
4013                )],
4014            )
4015            .unwrap()],
4016            |_| Ok(Some(format.clone())),
4017        )
4018        .unwrap()
4019        .try_into()
4020        .unwrap();
4021        let norm = ParameterGroupSpec::new(
4022            "norm",
4023            ParameterRole::Replicated,
4024            [ParameterMemberSpec::new(
4025                "model.layers.0.norm.weight",
4026                vec![64],
4027                MemberSharding::Replicated,
4028            )],
4029        )
4030        .unwrap();
4031        let owner = ParameterGroupOwner::execution_unit(layout.group_id(0).unwrap().clone(), 0);
4032        let description = ArchitectureParameterDescription::new(
4033            &graph,
4034            &layout,
4035            [physical.clone(), norm.clone()],
4036            [
4037                OwnedParameterGroupSpec::new(owner.clone(), physical.clone()),
4038                OwnedParameterGroupSpec::new(owner.clone(), norm.clone()),
4039            ],
4040        )
4041        .unwrap();
4042        let ownership =
4043            PartitionOwnership::new(false, false, std::iter::empty::<String>()).unwrap();
4044        let state = selected.requirements().state_layout();
4045        let state_plan =
4046            ArchitectureStatePartitionPlan::new([ArchitectureStatePartitionRule::group_units(
4047                0,
4048                0..state.len(),
4049            )]);
4050        let partition = ArchitecturePartition::from_description(
4051            &description,
4052            [(layout.group_id(0).unwrap().as_str(), 0..1)],
4053            ownership.clone(),
4054            state,
4055            &state_plan,
4056            (),
4057            NoAuxiliaryBoundarySchema::new(64),
4058        )
4059        .unwrap();
4060        let tasks =
4061            partitioned_replicated_text_materialization_tasks(&selected, &description, &partition)
4062                .unwrap();
4063        let task = tasks
4064            .iter()
4065            .find(|task| task.name() == "model.layers.0.mlp.weight")
4066            .unwrap();
4067        assert_eq!(task.output_companions().len(), 2);
4068
4069        let members = physical.members();
4070        let primary = ParameterGroupSpec::new(
4071            "primary",
4072            ParameterRole::FeedForwardIntermediate,
4073            [members[0].clone()],
4074        )
4075        .unwrap();
4076        let companions = ParameterGroupSpec::new(
4077            "companions",
4078            ParameterRole::FeedForwardIntermediate,
4079            members[1..].to_vec(),
4080        )
4081        .unwrap();
4082        let malformed = ArchitectureParameterDescription::new(
4083            &graph,
4084            &layout,
4085            [primary.clone(), companions.clone(), norm.clone()],
4086            [
4087                OwnedParameterGroupSpec::new(owner.clone(), primary),
4088                OwnedParameterGroupSpec::new(owner.clone(), companions),
4089                OwnedParameterGroupSpec::new(owner, norm),
4090            ],
4091        )
4092        .unwrap();
4093        let malformed_partition = ArchitecturePartition::from_description(
4094            &malformed,
4095            [(layout.group_id(0).unwrap().as_str(), 0..1)],
4096            ownership,
4097            state,
4098            &state_plan,
4099            (),
4100            NoAuxiliaryBoundarySchema::new(64),
4101        )
4102        .unwrap();
4103        let error = partitioned_replicated_text_materialization_tasks(
4104            &selected,
4105            &malformed,
4106            &malformed_partition,
4107        )
4108        .unwrap_err();
4109        assert!(error
4110            .to_string()
4111            .contains("outside its atomic parameter group"));
4112    }
4113
4114    #[test]
4115    fn partitioned_tasks_require_one_canonical_or_admitted_alias_topology_target() {
4116        let mut requirements = requirements();
4117        requirements.parameters[0].aliases = vec!["architecture.mlp.weight".into()];
4118        let selected = select_replicated_text_realization(
4119            &requirements,
4120            &request(LayerWeightResidency::FullyResident),
4121            &capabilities(),
4122        )
4123        .unwrap();
4124        let graph = selected.requirements().execution_graph().clone();
4125        let layout = selected.requirements().execution_units().clone();
4126        let owner = ParameterGroupOwner::execution_unit(layout.group_id(0).unwrap().clone(), 0);
4127        let ownership =
4128            PartitionOwnership::new(false, false, std::iter::empty::<String>()).unwrap();
4129        let state = selected.requirements().state_layout();
4130        let state_plan =
4131            ArchitectureStatePartitionPlan::new([ArchitectureStatePartitionRule::group_units(
4132                0,
4133                0..state.len(),
4134            )]);
4135
4136        let project = |primary_targets: &[&str]| {
4137            let mut groups = primary_targets
4138                .iter()
4139                .enumerate()
4140                .map(|(index, target)| {
4141                    ParameterGroupSpec::new(
4142                        format!("mlp-{index}"),
4143                        ParameterRole::FeedForwardIntermediate,
4144                        [ParameterMemberSpec::new(
4145                            *target,
4146                            vec![64, 64],
4147                            MemberSharding::Replicated,
4148                        )],
4149                    )
4150                    .unwrap()
4151                })
4152                .collect::<Vec<_>>();
4153            groups.push(
4154                ParameterGroupSpec::new(
4155                    "norm",
4156                    ParameterRole::Replicated,
4157                    [ParameterMemberSpec::new(
4158                        "model.layers.0.norm.weight",
4159                        vec![64],
4160                        MemberSharding::Replicated,
4161                    )],
4162                )
4163                .unwrap(),
4164            );
4165            let description = ArchitectureParameterDescription::new(
4166                &graph,
4167                &layout,
4168                groups.clone(),
4169                groups
4170                    .into_iter()
4171                    .map(|group| OwnedParameterGroupSpec::new(owner.clone(), group)),
4172            )
4173            .unwrap();
4174            let partition = ArchitecturePartition::from_description(
4175                &description,
4176                [(layout.group_id(0).unwrap().as_str(), 0..1)],
4177                ownership.clone(),
4178                state,
4179                &state_plan,
4180                (),
4181                NoAuxiliaryBoundarySchema::new(64),
4182            )
4183            .unwrap();
4184            partitioned_replicated_text_materialization_tasks(&selected, &description, &partition)
4185        };
4186
4187        let canonical = project(&["model.layers.0.mlp.weight"]).unwrap();
4188        assert!(canonical
4189            .iter()
4190            .any(|task| task.name() == "model.layers.0.mlp.weight"));
4191
4192        let aliased = project(&["architecture.mlp.weight"]).unwrap();
4193        let task = aliased
4194            .iter()
4195            .find(|task| task.name() == "model.layers.0.mlp.weight")
4196            .unwrap();
4197        assert_eq!(task.aliases(), ["architecture.mlp.weight"]);
4198
4199        let error = project(&["model.layers.0.mlp.weight", "architecture.mlp.weight"]).unwrap_err();
4200        assert!(error
4201            .to_string()
4202            .contains("resolves to 2 architecture topology targets"));
4203    }
4204
4205    #[test]
4206    fn selection_is_deterministic_and_keeps_source_format_distinct() {
4207        let disk = DenseDiskStreamLoadOptions::new(1234, 5678, 3, 2).unwrap();
4208        let request = request(LayerWeightResidency::DenseDiskStream(disk)).with_quantization(
4209            QuantizationRequest::Affine {
4210                group_size: 64,
4211                bits: 4,
4212            },
4213        );
4214        let left =
4215            select_replicated_text_realization(&requirements(), &request, &capabilities()).unwrap();
4216        let right =
4217            select_replicated_text_realization(&requirements(), &request, &capabilities()).unwrap();
4218        assert_eq!(left, right);
4219        assert_eq!(
4220            left.residency(),
4221            LayerWeightResidency::DenseDiskStream(disk)
4222        );
4223        assert_eq!(left.state().policy(), &paged_state());
4224        assert_eq!(left.state().layout(), requirements().state_layout());
4225        assert_eq!(left.parameters().len(), 2);
4226        assert_eq!(requirements().parameters().len(), 3);
4227        assert!(matches!(
4228            requirements().parameters()[1].presence(),
4229            ReplicatedTextParameterPresence::OptionalAbsent
4230        ));
4231        assert!(matches!(
4232            requirements().parameters()[2].role(),
4233            ReplicatedTextParameterRole::Normalization
4234        ));
4235        assert_eq!(requirements().parameters()[2].logical_shape(), [64]);
4236        assert_eq!(
4237            requirements().parameters()[2].transform_constraint(),
4238            ParameterTransformConstraint::None
4239        );
4240        assert_eq!(
4241            left.parameters()[0].lowering(),
4242            WeightLoweringKind::Transform
4243        );
4244        assert_ne!(
4245            format!("{:?}", left.parameters()[0].source_encoding()),
4246            format!("{:?}", left.parameters()[0].executable())
4247        );
4248    }
4249
4250    #[test]
4251    fn exact_tasks_are_the_authority_for_direct_derived_and_transform_sources() {
4252        use eredu_checkpoint::recipe::{DerivedWeightRecipe, RecipeDtype, RecipeMetadata};
4253
4254        let direct = select_replicated_text_realization(
4255            &requirements(),
4256            &request(LayerWeightResidency::FullyResident),
4257            &capabilities(),
4258        )
4259        .unwrap();
4260        let direct_tasks = replicated_text_materialization_tasks(&direct).unwrap();
4261        assert_eq!(
4262            direct_tasks[0].source_recipe().unwrap(),
4263            DerivedWeightRecipe::source(
4264                "blk.0.ffn.weight",
4265                eredu_checkpoint::store::TensorSelection::Full,
4266            )
4267        );
4268
4269        let recipe = DerivedWeightRecipe::source(
4270            "blk.0.ffn.weight",
4271            eredu_checkpoint::store::TensorSelection::Full,
4272        );
4273        let outputs = BTreeMap::from([(
4274            "model.layers.0.mlp.weight".into(),
4275            RecipeMetadata {
4276                shape: vec![64, 64],
4277                dtype: RecipeDtype::F16,
4278                byte_len: 64 * 64 * 2,
4279            },
4280        )]);
4281        let derived_requirements = requirements()
4282            .with_derived_recipes(
4283                BTreeMap::from([("model.layers.0.mlp.weight".into(), recipe.clone())]),
4284                outputs,
4285            )
4286            .unwrap();
4287        let derived = select_replicated_text_realization(
4288            &derived_requirements,
4289            &request(LayerWeightResidency::FullyResident),
4290            &capabilities(),
4291        )
4292        .unwrap();
4293        let derived_tasks = replicated_text_materialization_tasks(&derived).unwrap();
4294        assert_eq!(derived_tasks[0].lowering(), WeightLoweringKind::Derived);
4295        assert_eq!(derived_tasks[0].source_recipe().unwrap(), recipe);
4296
4297        let transformed = select_replicated_text_realization(
4298            &derived_requirements,
4299            &request(LayerWeightResidency::FullyResident).with_quantization(
4300                QuantizationRequest::Affine {
4301                    group_size: 64,
4302                    bits: 4,
4303                },
4304            ),
4305            &capabilities(),
4306        )
4307        .unwrap();
4308        let transformed_tasks = replicated_text_materialization_tasks(&transformed).unwrap();
4309        assert_eq!(
4310            transformed_tasks[0].lowering(),
4311            WeightLoweringKind::DerivedTransform
4312        );
4313        assert_eq!(transformed_tasks[0].source_recipe().unwrap(), recipe);
4314
4315        // These corruptions fail while projecting the cold exact plan; no
4316        // backend mechanism or checkpoint payload is available to perform work.
4317        let mut corrupt_direct = direct_tasks[0].clone();
4318        corrupt_direct.sources.push("unselected.weight".into());
4319        assert!(corrupt_direct.source_recipe().is_err());
4320        let mut corrupt_kind = derived_tasks[0].clone();
4321        corrupt_kind.lowering = WeightLoweringKind::Direct;
4322        assert!(corrupt_kind.source_recipe().is_err());
4323        let mut corrupt_recipe = derived_tasks[0].clone();
4324        corrupt_recipe.derived_recipe = Some(DerivedWeightRecipe::source(
4325            "unselected.weight",
4326            eredu_checkpoint::store::TensorSelection::Full,
4327        ));
4328        assert!(corrupt_recipe.source_recipe().is_err());
4329
4330        let member_output = RecipeMetadata {
4331            shape: vec![64, 64],
4332            dtype: RecipeDtype::F16,
4333            byte_len: 64 * 64 * 2,
4334        };
4335        let member_recipe = direct_tasks[0].source_recipe().unwrap();
4336        let project = |task, selected_bytes| {
4337            crate::AddressableBankParameter::new(
4338                "weight",
4339                task,
4340                member_recipe.clone(),
4341                member_output.clone(),
4342                selected_bytes,
4343                None,
4344            )
4345        };
4346        assert!(project(direct_tasks[0].clone(), member_output.byte_len()).is_ok());
4347        assert!(matches!(
4348            project(direct_tasks[0].clone(), member_output.byte_len() - 1),
4349            Err(crate::AddressableBankMemberError::SelectedByteMismatch { .. })
4350        ));
4351
4352        let mut corrupt_source = direct_tasks[0].clone();
4353        corrupt_source.source_encoding = SourceTensorEncoding::Safetensors(StoredDtype::F32);
4354        assert!(project(corrupt_source, member_output.byte_len()).is_err());
4355        let mut corrupt_executable = direct_tasks[0].clone();
4356        corrupt_executable.executable = LinearFormat::MxFp4;
4357        assert!(project(corrupt_executable, member_output.byte_len()).is_err());
4358        let mut corrupt_lowering = direct_tasks[0].clone();
4359        corrupt_lowering.lowering = WeightLoweringKind::Derived;
4360        assert!(project(corrupt_lowering, member_output.byte_len()).is_err());
4361
4362        let mut exact_transform = transformed_tasks[0].clone();
4363        let companion_owner = crate::ParameterGroupOwner::ExecutionUnit {
4364            group: crate::ExecutionGroupId::new("decoder").unwrap(),
4365            global_unit: 0,
4366        };
4367        exact_transform
4368            .set_output_companions(vec![
4369                ReplicatedTextOutputCompanion::new(
4370                    "model.layers.0.mlp.weight_scales",
4371                    eredu_nn::LinearCompanionRole::Scale,
4372                    vec![64, 1],
4373                    companion_owner.clone(),
4374                )
4375                .unwrap(),
4376                ReplicatedTextOutputCompanion::new(
4377                    "model.layers.0.mlp.weight_biases",
4378                    eredu_nn::LinearCompanionRole::AffineBias,
4379                    vec![64, 1],
4380                    companion_owner,
4381                )
4382                .unwrap(),
4383            ])
4384            .unwrap();
4385        let transformed_bytes =
4386            crate::selected_addressable_parameter_bytes(&exact_transform, &member_output).unwrap();
4387        assert!(crate::AddressableBankParameter::new(
4388            "weight",
4389            exact_transform.clone(),
4390            recipe.clone(),
4391            member_output.clone(),
4392            transformed_bytes,
4393            Some(
4394                crate::QuantizationCompanionBindings::new(
4395                    "weight_scales",
4396                    Some("weight_biases".into()),
4397                )
4398                .unwrap(),
4399            ),
4400        )
4401        .is_ok());
4402        assert!(crate::AddressableBankParameter::new(
4403            "weight",
4404            exact_transform,
4405            recipe.clone(),
4406            member_output.clone(),
4407            transformed_bytes,
4408            Some(
4409                crate::QuantizationCompanionBindings::new(
4410                    "drifted_scales",
4411                    Some("weight_biases".into()),
4412                )
4413                .unwrap(),
4414            ),
4415        )
4416        .is_err());
4417
4418        let mut scale_only = transformed_tasks[0].clone();
4419        scale_only.executable = LinearFormat::MxFp4;
4420        scale_only.lowering_descriptor = WeightLoweringDescriptor::new(
4421            scale_only.source_encoding.clone(),
4422            LinearFormat::MxFp4,
4423            scale_only.physical_shape.clone(),
4424            scale_only.logical_shape.clone(),
4425            scale_only.logical_shape.len().checked_sub(1),
4426        )
4427        .unwrap();
4428        scale_only
4429            .set_output_companions(vec![ReplicatedTextOutputCompanion::new(
4430                "model.layers.0.mlp.weight_scales",
4431                eredu_nn::LinearCompanionRole::Scale,
4432                vec![64, 2],
4433                crate::ParameterGroupOwner::ExecutionUnit {
4434                    group: crate::ExecutionGroupId::new("decoder").unwrap(),
4435                    global_unit: 0,
4436                },
4437            )
4438            .unwrap()])
4439            .unwrap();
4440        let scale_only_bytes =
4441            crate::selected_addressable_parameter_bytes(&scale_only, &member_output).unwrap();
4442        let scale_companions =
4443            crate::QuantizationCompanionBindings::new("weight_scales", None).unwrap();
4444        assert!(crate::AddressableBankParameter::new(
4445            "weight",
4446            scale_only.clone(),
4447            recipe.clone(),
4448            member_output.clone(),
4449            scale_only_bytes,
4450            Some(scale_companions),
4451        )
4452        .is_ok());
4453        let invented_bias = crate::QuantizationCompanionBindings::new(
4454            "weight_scales",
4455            Some("invented_bias".into()),
4456        )
4457        .unwrap();
4458        assert!(crate::AddressableBankParameter::new(
4459            "weight",
4460            scale_only,
4461            recipe,
4462            member_output,
4463            scale_only_bytes,
4464            Some(invented_bias),
4465        )
4466        .is_err());
4467    }
4468
4469    #[test]
4470    fn selection_reports_all_missing_mechanisms_together() {
4471        let capabilities = BackendMechanismCapabilities::new(
4472            NeuralOperatorCapabilities::NONE,
4473            Vec::new(),
4474            Vec::new(),
4475            StateMechanismCapabilities::new(Vec::new()),
4476        );
4477        let error = select_replicated_text_realization(
4478            &requirements(),
4479            &request(LayerWeightResidency::LayerwiseHost(
4480                LayerwiseLoadOptions::default(),
4481            )),
4482            &capabilities,
4483        )
4484        .unwrap_err();
4485        assert!(error.issues().len() >= 7, "{:?}", error.issues());
4486        assert!(error.issues().iter().any(|issue| issue.contains("exp")));
4487        assert!(error
4488            .issues()
4489            .iter()
4490            .any(|issue| issue.contains("weight lowering")));
4491    }
4492
4493    #[test]
4494    fn selection_rejects_paged_fixed_component_placement_even_when_reported() {
4495        let fixed = StateTensorPolicy::new(
4496            StateTensorRole::Recurrent,
4497            vec![
4498                StateTensorDimension::Batch,
4499                StateTensorDimension::fixed(8).unwrap(),
4500            ],
4501            StateTensorDtype::Float32,
4502            MutableStateResidency::LayerScopedOffloadable,
4503        )
4504        .unwrap();
4505        let mut requirements = requirements();
4506        requirements.state_layout = StateLayout::new(
4507            LayerSchedule::new(
4508                1,
4509                vec![LayerCachePolicy::key_value_with_fixed_state(
4510                    AttentionPolicy::Full,
4511                    1,
4512                    8,
4513                    vec![fixed],
4514                )
4515                .unwrap()],
4516            )
4517            .unwrap(),
4518        )
4519        .unwrap();
4520        requirements.state_access = ReplicatedTextStateAccess::AttentionWithFixed;
4521        let mut capabilities = capabilities();
4522        capabilities.state.components = (0..requirements.state_layout.len())
4523            .flat_map(|layer| {
4524                requirements
4525                    .state_layout
4526                    .components(layer)
4527                    .unwrap()
4528                    .iter()
4529                    .cloned()
4530                    .map(move |component| {
4531                        StateComponentMechanism::new(
4532                            layer,
4533                            component,
4534                            Some(StateComponentPlacement::Device),
4535                            Some(StateComponentPlacement::Paged),
4536                        )
4537                    })
4538            })
4539            .collect();
4540
4541        let error = select_replicated_text_realization(
4542            &requirements,
4543            &request(LayerWeightResidency::FullyResident),
4544            &capabilities,
4545        )
4546        .unwrap_err();
4547        assert!(error
4548            .issues()
4549            .iter()
4550            .any(|issue| issue.contains("incompatible Paged placement")));
4551    }
4552
4553    #[test]
4554    fn requirements_reject_state_layout_and_access_profile_mismatch() {
4555        let fixed = StateTensorPolicy::new(
4556            StateTensorRole::Recurrent,
4557            vec![
4558                StateTensorDimension::Batch,
4559                StateTensorDimension::fixed(8).unwrap(),
4560            ],
4561            StateTensorDtype::Float32,
4562            MutableStateResidency::LayerScopedOffloadable,
4563        )
4564        .unwrap();
4565        let layout = StateLayout::new(
4566            LayerSchedule::new(
4567                1,
4568                vec![LayerCachePolicy::key_value_with_fixed_state(
4569                    AttentionPolicy::Full,
4570                    1,
4571                    8,
4572                    vec![fixed],
4573                )
4574                .unwrap()],
4575            )
4576            .unwrap(),
4577        )
4578        .unwrap();
4579        let base = requirements();
4580        let error = ReplicatedTextRequirements::new(
4581            base.architecture_identity,
4582            base.operators,
4583            base.execution_graph,
4584            base.execution_units,
4585            base.group_transports,
4586            layout,
4587            ReplicatedTextStateAccess::KeyValue,
4588            base.parameters,
4589        )
4590        .unwrap_err();
4591        assert!(error.message().contains("does not match component roles"));
4592    }
4593
4594    #[test]
4595    fn transform_selection_rejects_incompatible_exact_geometry() {
4596        for quantization in [
4597            QuantizationRequest::Affine {
4598                group_size: 96,
4599                bits: 4,
4600            },
4601            QuantizationRequest::Affine {
4602                group_size: 256,
4603                bits: 4,
4604            },
4605            QuantizationRequest::Affine {
4606                group_size: 0,
4607                bits: 4,
4608            },
4609            QuantizationRequest::Affine {
4610                group_size: u32::MAX,
4611                bits: 4,
4612            },
4613            QuantizationRequest::Affine {
4614                group_size: 32,
4615                bits: 0,
4616            },
4617            QuantizationRequest::Affine {
4618                group_size: 32,
4619                bits: 7,
4620            },
4621        ] {
4622            let error = select_replicated_text_realization(
4623                &requirements(),
4624                &request(LayerWeightResidency::FullyResident).with_quantization(quantization),
4625                &capabilities(),
4626            )
4627            .unwrap_err();
4628            assert!(error
4629                .issues()
4630                .iter()
4631                .any(|issue| issue.contains("invalid replicated text contract")));
4632        }
4633
4634        let mut indivisible = requirements();
4635        indivisible.parameters[0].logical_shape = vec![64, 48];
4636        let error = select_replicated_text_realization(
4637            &indivisible,
4638            &request(LayerWeightResidency::FullyResident)
4639                .with_quantization(QuantizationRequest::MxFp4),
4640            &capabilities(),
4641        )
4642        .unwrap_err();
4643        assert!(error
4644            .issues()
4645            .iter()
4646            .any(|issue| issue.contains("MXFP4 packed extent 48")));
4647    }
4648
4649    #[test]
4650    fn exact_source_and_physical_geometry_fail_before_construction_or_payload() {
4651        for mutate in [
4652            |requirement: &mut ReplicatedTextParameterRequirement| {
4653                requirement.source_encoding =
4654                    Some(SourceTensorEncoding::Safetensors(StoredDtype::U8));
4655            },
4656            |requirement: &mut ReplicatedTextParameterRequirement| {
4657                requirement.physical_shape = Some(vec![64, 32]);
4658            },
4659        ] {
4660            let mut requirements = requirements();
4661            mutate(&mut requirements.parameters[0]);
4662            let selected = select_replicated_text_realization(
4663                &requirements,
4664                &request(LayerWeightResidency::FullyResident),
4665                &capabilities(),
4666            );
4667            let error = selected.unwrap_err();
4668            assert!(error
4669                .issues()
4670                .iter()
4671                .any(|issue| issue.contains("weight lowering")));
4672        }
4673    }
4674
4675    #[test]
4676    fn missing_tensor_parallel_grouped_partial_fails_before_construction_or_forward() {
4677        let requirements = requirements().with_grouped_operations([
4678            GroupedOperationRequirement::GatedProduct,
4679            GroupedOperationRequirement::GatedProductTensorParallelPartial,
4680        ]);
4681        let capabilities =
4682            capabilities().with_grouped_operations([GroupedOperationRequirement::GatedProduct]);
4683        let selected = select_replicated_text_realization(
4684            &requirements,
4685            &request(LayerWeightResidency::FullyResident),
4686            &capabilities,
4687        );
4688        let error = selected.unwrap_err();
4689        assert!(error
4690            .issues()
4691            .iter()
4692            .any(|issue| { issue.contains("GatedProductTensorParallelPartial") }));
4693    }
4694}