eredu_runtime/
replicated_text.rs

1//! Selection contracts for replicated text architectures.
2
3use std::{
4    collections::BTreeSet,
5    path::{Path, PathBuf},
6};
7
8use eredu_checkpoint::{LinearFormat, SourceTensorEncoding};
9use eredu_core::{ParallelTopology, QuantizationRequest, SessionCapabilities};
10use eredu_nn::{NeuralBackend, NeuralOperatorCapabilities};
11
12use crate::{
13    ArchitectureGroupTransport, CacheResidencyPolicy, ExecutionGraph, ExecutionUnitLayout,
14    LayerWeightResidency, LayeredArchitecture, RuntimeState, StateLayout,
15};
16
17/// Statically dispatched text-input seam for an ordinary layered decoder.
18///
19/// Hybrid, routed, composite, partitioned, prediction, and realtime execution
20/// use separate extension contracts rather than adding requirements here.
21pub trait ReplicatedTextArchitecture<B, S>: LayeredArchitecture<B, S>
22where
23    B: NeuralBackend,
24    S: RuntimeState<B>,
25{
26    /// Forms the architecture-owned borrowed input for one text pass.
27    fn text_input<'a>(tokens: &'a B::Tensor, mask: Option<&'a B::Tensor>) -> Self::Input<'a>;
28}
29
30/// Backend implementation route for one source-to-executable weight lowering.
31#[derive(Debug, Clone, Copy, Eq, PartialEq)]
32#[non_exhaustive]
33pub enum WeightLoweringKind {
34    /// The admitted source encoding is retained by the executable operator.
35    Direct,
36    /// Payload materialization performs an admitted transformation.
37    Transform,
38}
39
40/// One exact weight lowering implemented by a backend.
41#[derive(Debug, Clone, Eq, PartialEq)]
42pub struct WeightLoweringCapability {
43    /// Exact neutral lowering request implemented by this capability.
44    descriptor: WeightLoweringDescriptor,
45    /// Whether the lowering is direct or transforming.
46    kind: WeightLoweringKind,
47}
48
49impl WeightLoweringCapability {
50    /// Creates one exact backend lowering mechanism.
51    pub fn new(descriptor: WeightLoweringDescriptor, kind: WeightLoweringKind) -> Self {
52        Self { descriptor, kind }
53    }
54
55    /// Returns the admitted source encoding.
56    pub const fn source(&self) -> &SourceTensorEncoding {
57        self.descriptor.source()
58    }
59
60    /// Returns the executable format produced by this mechanism.
61    pub const fn executable(&self) -> LinearFormat {
62        self.descriptor.executable()
63    }
64
65    /// Returns whether materialization retains or transforms the source.
66    pub const fn kind(&self) -> WeightLoweringKind {
67        self.kind
68    }
69
70    /// Returns the exact geometry-bearing lowering request.
71    pub const fn descriptor(&self) -> &WeightLoweringDescriptor {
72        &self.descriptor
73    }
74}
75
76/// Exact source-to-executable lowering query presented to a backend.
77#[derive(Debug, Clone, Eq, PartialEq)]
78pub struct WeightLoweringDescriptor {
79    source: SourceTensorEncoding,
80    executable: LinearFormat,
81    physical_shape: Vec<usize>,
82    logical_shape: Vec<usize>,
83    packed_axis: Option<usize>,
84}
85
86impl WeightLoweringDescriptor {
87    /// Creates a geometry-bearing lowering query.
88    pub fn new(
89        source: SourceTensorEncoding,
90        executable: LinearFormat,
91        physical_shape: Vec<usize>,
92        logical_shape: Vec<usize>,
93        packed_axis: Option<usize>,
94    ) -> Result<Self, ReplicatedTextContractError> {
95        if physical_shape.is_empty()
96            || physical_shape.contains(&0)
97            || logical_shape.is_empty()
98            || logical_shape.contains(&0)
99            || physical_shape.len() != logical_shape.len()
100        {
101            return Err(ReplicatedTextContractError::invalid(
102                "weight lowering requires positive physical and logical shapes of equal rank",
103            ));
104        }
105        if packed_axis.is_some_and(|axis| axis >= logical_shape.len()) {
106            return Err(ReplicatedTextContractError::invalid(
107                "weight lowering packed axis is outside the logical shape",
108            ));
109        }
110        Ok(Self {
111            source,
112            executable,
113            physical_shape,
114            logical_shape,
115            packed_axis,
116        })
117    }
118
119    /// Returns the admitted source encoding.
120    pub const fn source(&self) -> &SourceTensorEncoding {
121        &self.source
122    }
123
124    /// Returns the selected executable format.
125    pub const fn executable(&self) -> LinearFormat {
126        self.executable
127    }
128
129    /// Returns the admitted physical source shape.
130    pub fn physical_shape(&self) -> &[usize] {
131        &self.physical_shape
132    }
133
134    /// Returns the architecture-declared logical shape.
135    pub fn logical_shape(&self) -> &[usize] {
136        &self.logical_shape
137    }
138
139    /// Returns the executable packing axis, when the parameter is packable.
140    pub const fn packed_axis(&self) -> Option<usize> {
141        self.packed_axis
142    }
143
144    /// Returns the exact extent along the packing axis.
145    pub fn packed_extent(&self) -> Option<usize> {
146        self.packed_axis.map(|axis| self.logical_shape[axis])
147    }
148}
149
150/// Weight-residency mechanism implemented by a backend.
151#[derive(Debug, Clone, Copy, Eq, PartialEq)]
152#[non_exhaustive]
153pub enum WeightResidencyMechanism {
154    /// All parameters remain device resident.
155    Resident,
156    /// A bounded device window is staged from host storage.
157    Windowed,
158    /// Bounded host and device windows are populated from disk.
159    DiskStreamed,
160}
161
162/// Mutable-state residency mechanism implemented by a backend.
163#[derive(Debug, Clone, Copy, Eq, PartialEq)]
164#[non_exhaustive]
165pub enum StateResidencyMechanism {
166    /// State tensors remain on the execution device.
167    Device,
168    /// State blocks use bounded paged storage.
169    Paged,
170}
171
172/// Architecture-valid transform target for one linear parameter.
173#[derive(Debug, Clone, Eq, PartialEq)]
174pub struct ParameterTransformTarget {
175    /// Requested load-time transform.
176    request: QuantizationRequest,
177    /// Executable format produced for this parameter.
178    executable: LinearFormat,
179    /// Exact geometry that the backend lowering must accept.
180    descriptor: WeightLoweringDescriptor,
181}
182
183impl ParameterTransformTarget {
184    /// Creates one architecture-admitted load-time transform target.
185    fn new(
186        request: QuantizationRequest,
187        executable: LinearFormat,
188        descriptor: WeightLoweringDescriptor,
189    ) -> Self {
190        Self {
191            request,
192            executable,
193            descriptor,
194        }
195    }
196
197    /// Returns the caller request selecting this transform.
198    pub const fn request(&self) -> QuantizationRequest {
199        self.request
200    }
201
202    /// Returns the architecture-admitted executable format.
203    pub const fn executable(&self) -> LinearFormat {
204        self.executable
205    }
206
207    /// Returns the exact neutral lowering query.
208    pub const fn descriptor(&self) -> &WeightLoweringDescriptor {
209        &self.descriptor
210    }
211}
212
213/// Architecture declaration of whether and how a parameter may be transformed.
214#[derive(Debug, Clone, Copy, Eq, PartialEq)]
215#[non_exhaustive]
216pub enum ParameterTransformConstraint {
217    /// This parameter is not an executable affine projection weight.
218    None,
219    /// The declared axis is the input/packing axis of a linear parameter.
220    Linear {
221        /// Axis whose extent is grouped or blocked by executable packing.
222        packed_axis: usize,
223    },
224}
225
226/// Architecture-owned semantic role of one logical parameter.
227#[derive(Debug, Clone, Copy, Eq, PartialEq)]
228#[non_exhaustive]
229pub enum ReplicatedTextParameterRole {
230    /// Token lookup table.
231    Embedding,
232    /// Executable affine projection weight.
233    LinearWeight,
234    /// Learned affine projection bias.
235    LinearBias,
236    /// Learned normalization scale or offset.
237    Normalization,
238    /// Physical scale, zero-point, or packed-format companion.
239    FormatCompanion,
240    /// Another architecture-declared non-linear parameter.
241    Other,
242}
243
244/// Architecture-owned location of one replicated-text parameter.
245#[derive(Debug, Clone, Eq, PartialEq)]
246#[non_exhaustive]
247pub enum ReplicatedTextParameterOwner {
248    /// Pinned module selected by a stable architecture role.
249    StaticRole(String),
250    /// One architecture-global execution unit.
251    ExecutionUnit {
252        /// Stable execution-group identity.
253        group: String,
254        /// Group-local architecture-global unit index.
255        unit: usize,
256    },
257}
258
259/// Exact admitted presence or derivation of one logical parameter.
260#[derive(Debug, Clone, Eq, PartialEq)]
261#[non_exhaustive]
262pub enum ReplicatedTextParameterPresence {
263    /// A required physical source was selected.
264    Required,
265    /// An optional physical source was present and selected.
266    OptionalPresent,
267    /// An optional architecture parameter is absent from this artifact.
268    OptionalAbsent,
269    /// The value is tied to another canonical logical parameter.
270    Tied {
271        /// Canonical identity supplying the value.
272        target: String,
273    },
274    /// The value is produced by an architecture-owned recipe.
275    Derived {
276        /// Stable recipe identity.
277        recipe: String,
278    },
279}
280
281impl ReplicatedTextParameterPresence {
282    /// Returns whether selection must choose a backend lowering.
283    pub fn has_physical_source(&self) -> bool {
284        matches!(self, Self::Required | Self::OptionalPresent)
285    }
286}
287
288/// Exact admitted source and executable constraints for one logical parameter.
289#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
290pub struct ReplicatedTextPhysicalSource {
291    tensor: String,
292    shard: PathBuf,
293    output: String,
294}
295
296impl ReplicatedTextPhysicalSource {
297    /// Records one exact physical tensor, admitted shard, and selected output.
298    pub fn new(
299        tensor: impl Into<String>,
300        shard: impl Into<PathBuf>,
301        output: impl Into<String>,
302    ) -> Result<Self, ReplicatedTextContractError> {
303        let tensor = tensor.into();
304        let shard = shard.into();
305        let output = output.into();
306        if tensor.trim().is_empty() || shard.as_os_str().is_empty() || output.trim().is_empty() {
307            return Err(ReplicatedTextContractError::invalid(
308                "physical source tensor, shard, and output must be non-empty",
309            ));
310        }
311        Ok(Self {
312            tensor,
313            shard,
314            output,
315        })
316    }
317
318    /// Physical tensor identity in the admitted container.
319    pub fn tensor(&self) -> &str {
320        &self.tensor
321    }
322    /// Canonical admitted payload shard.
323    pub fn shard(&self) -> &Path {
324        &self.shard
325    }
326    /// Exact logical output selected from the physical tensor.
327    pub fn output(&self) -> &str {
328        &self.output
329    }
330}
331
332/// Exact admitted source and executable constraints for one logical parameter.
333#[derive(Debug, Clone, Eq, PartialEq)]
334pub struct ReplicatedTextParameterRequirement {
335    /// Canonical logical parameter identity.
336    name: String,
337    /// Physical outputs admitted as sources for this logical parameter.
338    sources: Vec<String>,
339    /// Exact shard and multi-output provenance for the physical input.
340    physical_sources: Vec<ReplicatedTextPhysicalSource>,
341    /// All admitted aliases for the logical parameter.
342    aliases: Vec<String>,
343    /// Encoding of the selected physical source, when present.
344    source_encoding: Option<SourceTensorEncoding>,
345    /// Exact selected physical source shape, when present.
346    physical_shape: Option<Vec<usize>>,
347    /// Architecture-declared logical tensor shape.
348    logical_shape: Vec<usize>,
349    /// Architecture-owned semantic parameter role.
350    role: ReplicatedTextParameterRole,
351    /// Architecture-owned static/group/unit location.
352    owner: ReplicatedTextParameterOwner,
353    /// Exact artifact presence, tie, or derivation.
354    presence: ReplicatedTextParameterPresence,
355    /// Architecture-selected native executable format.
356    native_executable: LinearFormat,
357    /// Exact architecture-owned transform eligibility and packing axis.
358    transform: ParameterTransformConstraint,
359}
360
361impl ReplicatedTextParameterRequirement {
362    /// Creates one exact logical-parameter requirement.
363    #[allow(
364        clippy::too_many_arguments,
365        reason = "the constructor validates one complete immutable catalog record"
366    )]
367    pub fn new(
368        name: impl Into<String>,
369        sources: Vec<String>,
370        physical_sources: Vec<ReplicatedTextPhysicalSource>,
371        aliases: Vec<String>,
372        source_encoding: Option<SourceTensorEncoding>,
373        physical_shape: Option<Vec<usize>>,
374        logical_shape: Vec<usize>,
375        native_executable: LinearFormat,
376        role: ReplicatedTextParameterRole,
377        owner: ReplicatedTextParameterOwner,
378        presence: ReplicatedTextParameterPresence,
379        transform: ParameterTransformConstraint,
380    ) -> Result<Self, ReplicatedTextContractError> {
381        let name = name.into();
382        if name.trim().is_empty() {
383            return Err(ReplicatedTextContractError::invalid(
384                "logical parameter identity is empty",
385            ));
386        }
387        if sources.iter().any(|source| source.trim().is_empty())
388            || aliases.iter().any(|alias| alias.trim().is_empty())
389        {
390            return Err(ReplicatedTextContractError::invalid(format!(
391                "logical parameter {name:?} has an empty physical identity"
392            )));
393        }
394        let has_source = presence.has_physical_source();
395        if has_source
396            != (!sources.is_empty() && source_encoding.is_some() && physical_shape.is_some())
397        {
398            return Err(ReplicatedTextContractError::invalid(format!(
399                "logical parameter {name:?} has inconsistent source presence"
400            )));
401        }
402        let provenance_required =
403            has_source || matches!(presence, ReplicatedTextParameterPresence::Derived { .. });
404        if provenance_required != !physical_sources.is_empty() {
405            return Err(ReplicatedTextContractError::invalid(format!(
406                "logical parameter {name:?} has inconsistent physical provenance"
407            )));
408        }
409        if has_source
410            && physical_sources
411                .iter()
412                .any(|source| !sources.iter().any(|name| name == source.tensor()))
413        {
414            return Err(ReplicatedTextContractError::invalid(format!(
415                "logical parameter {name:?} has provenance outside its selected sources"
416            )));
417        }
418        if physical_shape
419            .as_ref()
420            .is_some_and(|shape| shape.is_empty() || shape.contains(&0))
421        {
422            return Err(ReplicatedTextContractError::invalid(format!(
423                "logical parameter {name:?} has an invalid physical shape"
424            )));
425        }
426        if logical_shape.is_empty() || logical_shape.contains(&0) {
427            return Err(ReplicatedTextContractError::invalid(format!(
428                "logical parameter {name:?} has an invalid shape {logical_shape:?}"
429            )));
430        }
431        if let ParameterTransformConstraint::Linear { packed_axis } = transform {
432            if packed_axis >= logical_shape.len() {
433                return Err(ReplicatedTextContractError::invalid(format!(
434                    "logical parameter {name:?} has packing axis {packed_axis} outside shape {logical_shape:?}"
435                )));
436            }
437        }
438        Ok(Self {
439            name,
440            sources,
441            physical_sources,
442            aliases,
443            source_encoding,
444            physical_shape,
445            logical_shape,
446            role,
447            owner,
448            presence,
449            native_executable,
450            transform,
451        })
452    }
453
454    /// Returns the canonical logical identity.
455    pub fn name(&self) -> &str {
456        &self.name
457    }
458
459    /// Returns exact admitted physical source identities.
460    pub fn sources(&self) -> &[String] {
461        &self.sources
462    }
463
464    /// Returns exact admitted shard and multi-output provenance.
465    pub fn physical_sources(&self) -> &[ReplicatedTextPhysicalSource] {
466        &self.physical_sources
467    }
468
469    /// Returns all architecture-admitted alternative source identities.
470    pub fn aliases(&self) -> &[String] {
471        &self.aliases
472    }
473
474    /// Returns the admitted physical source encoding.
475    pub const fn source_encoding(&self) -> Option<&SourceTensorEncoding> {
476        self.source_encoding.as_ref()
477    }
478
479    /// Returns the selected physical shape, when a source is present.
480    pub fn physical_shape(&self) -> Option<&[usize]> {
481        self.physical_shape.as_deref()
482    }
483
484    /// Returns the architecture-declared logical shape.
485    pub fn logical_shape(&self) -> &[usize] {
486        &self.logical_shape
487    }
488
489    /// Returns the architecture-owned semantic role.
490    pub const fn role(&self) -> ReplicatedTextParameterRole {
491        self.role
492    }
493
494    /// Returns the architecture-owned static/group/unit location.
495    pub const fn owner(&self) -> &ReplicatedTextParameterOwner {
496        &self.owner
497    }
498
499    /// Returns exact artifact presence, tie, or derivation.
500    pub const fn presence(&self) -> &ReplicatedTextParameterPresence {
501        &self.presence
502    }
503
504    /// Returns exact transform eligibility and packing geometry.
505    pub const fn transform_constraint(&self) -> ParameterTransformConstraint {
506        self.transform
507    }
508
509    /// Returns the architecture-native executable format.
510    pub const fn native_executable(&self) -> LinearFormat {
511        self.native_executable
512    }
513
514    /// Resolves a caller transform through architecture-owned constraints.
515    pub fn transform_target(
516        &self,
517        request: QuantizationRequest,
518    ) -> Result<Option<ParameterTransformTarget>, ReplicatedTextContractError> {
519        let packed_axis = match self.transform {
520            ParameterTransformConstraint::None => return Ok(None),
521            ParameterTransformConstraint::Linear { packed_axis } => packed_axis,
522        };
523        let extent = self.logical_shape[packed_axis];
524        let executable = match request {
525            QuantizationRequest::Affine { group_size, bits } => {
526                let group_size = i32::try_from(group_size).map_err(|_| {
527                    ReplicatedTextContractError::invalid("affine group size exceeds i32")
528                })?;
529                let format = eredu_checkpoint::AffineQuantization::new(group_size, i32::from(bits))
530                    .map_err(|error| ReplicatedTextContractError::invalid(error.to_string()))?;
531                let group_size = usize::try_from(format.group_size).map_err(|_| {
532                    ReplicatedTextContractError::invalid("affine group size is negative")
533                })?;
534                if group_size > extent || !extent.is_multiple_of(group_size) {
535                    return Err(ReplicatedTextContractError::invalid(format!(
536                        "affine group size {group_size} does not divide packed extent {extent}"
537                    )));
538                }
539                LinearFormat::Affine(format)
540            }
541            QuantizationRequest::MxFp4 => {
542                const MXFP4_BLOCK_SIZE: usize = 32;
543                if !extent.is_multiple_of(MXFP4_BLOCK_SIZE) {
544                    return Err(ReplicatedTextContractError::invalid(format!(
545                        "MXFP4 packed extent {extent} is not divisible by block size {MXFP4_BLOCK_SIZE}"
546                    )));
547                }
548                LinearFormat::MxFp4
549            }
550            _ => {
551                return Err(ReplicatedTextContractError::invalid(
552                    "unknown load-time transform request",
553                ))
554            }
555        };
556        let descriptor = self.lowering_descriptor(executable)?;
557        Ok(Some(ParameterTransformTarget::new(
558            request, executable, descriptor,
559        )))
560    }
561
562    /// Forms the exact backend lowering query for one admitted executable format.
563    pub fn lowering_descriptor(
564        &self,
565        executable: LinearFormat,
566    ) -> Result<WeightLoweringDescriptor, ReplicatedTextContractError> {
567        let packed_axis = match self.transform {
568            ParameterTransformConstraint::None => None,
569            ParameterTransformConstraint::Linear { packed_axis } => Some(packed_axis),
570        };
571        let packed_axis = packed_axis.or_else(|| {
572            (self.role == ReplicatedTextParameterRole::Embedding
573                && executable != LinearFormat::Dense)
574                .then(|| self.logical_shape.len() - 1)
575        });
576        WeightLoweringDescriptor::new(
577            self.source_encoding.clone().ok_or_else(|| {
578                ReplicatedTextContractError::invalid(format!(
579                    "logical parameter {:?} has no physical lowering source",
580                    self.name
581                ))
582            })?,
583            executable,
584            self.physical_shape.clone().ok_or_else(|| {
585                ReplicatedTextContractError::invalid(format!(
586                    "logical parameter {:?} has no physical source shape",
587                    self.name
588                ))
589            })?,
590            self.logical_shape.clone(),
591            packed_axis,
592        )
593    }
594}
595
596/// Invalid public replicated-text contract construction.
597#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
598#[error("invalid replicated text contract: {message}")]
599pub struct ReplicatedTextContractError {
600    message: String,
601}
602
603impl ReplicatedTextContractError {
604    fn invalid(message: impl Into<String>) -> Self {
605        Self {
606            message: message.into(),
607        }
608    }
609
610    /// Returns the stable semantic diagnostic.
611    pub fn message(&self) -> &str {
612        &self.message
613    }
614}
615
616/// Exact architecture and artifact requirements for replicated text execution.
617#[derive(Debug, Clone, Eq, PartialEq)]
618pub struct ReplicatedTextRequirements {
619    /// Optional neural operations required by the architecture equations.
620    operators: NeuralOperatorCapabilities,
621    /// Stable architecture-owned execution graph.
622    execution_graph: ExecutionGraph,
623    /// Exact group-major execution-unit geometry.
624    execution_units: ExecutionUnitLayout,
625    /// Architecture-owned transport semantics in graph-group order.
626    group_transports: Vec<ArchitectureGroupTransport>,
627    /// Complete architecture-owned mutable-state geometry.
628    state_layout: StateLayout,
629    /// Canonical logical parameter requirements.
630    parameters: Vec<ReplicatedTextParameterRequirement>,
631    grouped_operations: Vec<GroupedOperationRequirement>,
632}
633
634impl ReplicatedTextRequirements {
635    /// Creates exact requirements from architecture and admitted-artifact facts only.
636    pub fn new(
637        operators: NeuralOperatorCapabilities,
638        execution_graph: ExecutionGraph,
639        execution_units: ExecutionUnitLayout,
640        group_transports: Vec<ArchitectureGroupTransport>,
641        state_layout: StateLayout,
642        parameters: Vec<ReplicatedTextParameterRequirement>,
643    ) -> Result<Self, ReplicatedTextContractError> {
644        if group_transports.len() != execution_graph.groups().len() {
645            return Err(ReplicatedTextContractError::invalid(format!(
646                "{} group transports do not match {} execution groups",
647                group_transports.len(),
648                execution_graph.groups().len()
649            )));
650        }
651        let mut names = BTreeSet::new();
652        if parameters
653            .iter()
654            .any(|parameter| !names.insert(parameter.name()))
655        {
656            return Err(ReplicatedTextContractError::invalid(
657                "logical parameter identities are not unique",
658            ));
659        }
660        Ok(Self {
661            operators,
662            execution_graph,
663            execution_units,
664            group_transports,
665            state_layout,
666            parameters,
667            grouped_operations: Vec::new(),
668        })
669    }
670
671    /// Declares exact grouped operations required by this architecture path.
672    pub fn with_grouped_operations(
673        mut self,
674        operations: impl IntoIterator<Item = GroupedOperationRequirement>,
675    ) -> Self {
676        self.grouped_operations = operations.into_iter().collect();
677        self
678    }
679
680    /// Returns required optional neural-operation semantics.
681    pub const fn operators(&self) -> NeuralOperatorCapabilities {
682        self.operators
683    }
684    /// Returns the architecture-owned execution graph.
685    pub const fn execution_graph(&self) -> &ExecutionGraph {
686        &self.execution_graph
687    }
688    /// Returns group-major execution-unit geometry.
689    pub const fn execution_units(&self) -> &ExecutionUnitLayout {
690        &self.execution_units
691    }
692    /// Returns architecture-owned group transports.
693    pub fn group_transports(&self) -> &[ArchitectureGroupTransport] {
694        &self.group_transports
695    }
696    /// Returns complete mutable-state geometry.
697    pub const fn state_layout(&self) -> &StateLayout {
698        &self.state_layout
699    }
700    /// Returns canonical logical parameter requirements.
701    pub fn parameters(&self) -> &[ReplicatedTextParameterRequirement] {
702        &self.parameters
703    }
704    /// Returns exact grouped operations required before construction.
705    pub fn grouped_operations(&self) -> &[GroupedOperationRequirement] {
706        &self.grouped_operations
707    }
708}
709
710/// One required grouped-compute mechanism.
711#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
712#[non_exhaustive]
713pub enum GroupedOperationRequirement {
714    /// Ordinary grouped gated-product output.
715    GatedProduct,
716    /// Rank-local gated-product partial with an explicit post-reduce term.
717    GatedProductTensorParallelPartial,
718    /// Ordinary grouped ReLU-squared output.
719    Relu2,
720    /// Rank-local ReLU-squared partial with an explicit post-reduce term.
721    Relu2TensorParallelPartial,
722}
723
724/// Family- and execution-class-neutral backend mechanism report.
725#[derive(Debug, Clone, Eq, PartialEq)]
726pub struct BackendMechanismCapabilities {
727    /// Optional neural operations implemented by the backend.
728    operators: NeuralOperatorCapabilities,
729    /// Exact admitted source-to-executable lowerings.
730    weight_lowerings: Vec<WeightLoweringCapability>,
731    /// Ordinary parameter residency mechanisms.
732    weight_residencies: Vec<WeightResidencyMechanism>,
733    /// Mutable-state residency mechanisms.
734    state_residencies: Vec<StateResidencyMechanism>,
735    /// Exact session facilities implemented by the constructed session.
736    session: SessionCapabilities,
737    /// Prompt-cache persistence mechanism is available.
738    prompt_cache: bool,
739    /// Exact completion ownership is implemented for submitted work.
740    exact_completion: bool,
741    grouped_operations: Vec<GroupedOperationRequirement>,
742}
743
744impl BackendMechanismCapabilities {
745    /// Creates a fail-closed mechanism report.
746    pub fn new(
747        operators: NeuralOperatorCapabilities,
748        weight_lowerings: Vec<WeightLoweringCapability>,
749        weight_residencies: Vec<WeightResidencyMechanism>,
750        state_residencies: Vec<StateResidencyMechanism>,
751    ) -> Self {
752        Self {
753            operators,
754            weight_lowerings,
755            weight_residencies,
756            state_residencies,
757            session: SessionCapabilities::default(),
758            prompt_cache: false,
759            exact_completion: false,
760            grouped_operations: Vec::new(),
761        }
762    }
763
764    /// Adds supported session-observation and persistence mechanisms.
765    pub const fn with_session(mut self, session: SessionCapabilities) -> Self {
766        self.session = session;
767        self
768    }
769    /// Declares prompt-cache persistence support.
770    pub const fn with_prompt_cache(mut self, supported: bool) -> Self {
771        self.prompt_cache = supported;
772        self
773    }
774    /// Declares exact native-completion ownership support.
775    pub const fn with_exact_completion(mut self, supported: bool) -> Self {
776        self.exact_completion = supported;
777        self
778    }
779    /// Declares exact grouped operation mechanisms.
780    pub fn with_grouped_operations(
781        mut self,
782        operations: impl IntoIterator<Item = GroupedOperationRequirement>,
783    ) -> Self {
784        self.grouped_operations = operations.into_iter().collect();
785        self
786    }
787    /// Returns neural-operation mechanisms.
788    pub const fn operators(&self) -> NeuralOperatorCapabilities {
789        self.operators
790    }
791    /// Returns source-to-executable weight-lowering mechanisms.
792    pub fn weight_lowerings(&self) -> &[WeightLoweringCapability] {
793        &self.weight_lowerings
794    }
795    /// Returns weight-residency mechanisms.
796    pub fn weight_residencies(&self) -> &[WeightResidencyMechanism] {
797        &self.weight_residencies
798    }
799    /// Returns mutable-state residency mechanisms.
800    pub fn state_residencies(&self) -> &[StateResidencyMechanism] {
801        &self.state_residencies
802    }
803    /// Returns session-observation and persistence mechanisms.
804    pub const fn session(&self) -> SessionCapabilities {
805        self.session
806    }
807    /// Returns whether prompt-cache persistence is supported.
808    pub const fn prompt_cache(&self) -> bool {
809        self.prompt_cache
810    }
811    /// Returns whether exact native-completion ownership is supported.
812    pub const fn exact_completion(&self) -> bool {
813        self.exact_completion
814    }
815    /// Returns exact grouped operation mechanisms.
816    pub fn grouped_operations(&self) -> &[GroupedOperationRequirement] {
817        &self.grouped_operations
818    }
819}
820
821/// Caller choices resolved while selecting one replicated text realization.
822#[derive(Debug, Clone, Eq, PartialEq)]
823pub struct ReplicatedTextSelectionRequest {
824    /// Requested execution topology.
825    topology: Option<ParallelTopology>,
826    /// Requested ordinary parameter residency.
827    residency: LayerWeightResidency,
828    /// Requested mutable-state implementation and its exact residency policy.
829    state: CacheResidencyPolicy,
830    /// Optional load-time transform.
831    quantization: Option<QuantizationRequest>,
832    /// Requested optional session facilities.
833    session: SessionCapabilities,
834    /// Whether prompt-cache persistence is requested.
835    prompt_cache: bool,
836    /// Whether exact completion ownership is requested.
837    exact_completion: bool,
838}
839
840impl ReplicatedTextSelectionRequest {
841    /// Creates a replicated request with fail-closed optional facilities.
842    pub fn new(residency: LayerWeightResidency, state: CacheResidencyPolicy) -> Self {
843        Self {
844            topology: None,
845            residency,
846            state,
847            quantization: None,
848            session: SessionCapabilities::default(),
849            prompt_cache: false,
850            exact_completion: false,
851        }
852    }
853    /// Sets the requested topology.
854    pub const fn with_topology(mut self, topology: ParallelTopology) -> Self {
855        self.topology = Some(topology);
856        self
857    }
858    /// Sets the optional load-time transform.
859    pub const fn with_quantization(mut self, quantization: QuantizationRequest) -> Self {
860        self.quantization = Some(quantization);
861        self
862    }
863    /// Sets requested session facilities.
864    pub const fn with_session(mut self, session: SessionCapabilities) -> Self {
865        self.session = session;
866        self
867    }
868    /// Requests prompt-cache persistence.
869    pub const fn with_prompt_cache(mut self, required: bool) -> Self {
870        self.prompt_cache = required;
871        self
872    }
873    /// Requests exact completion ownership.
874    pub const fn with_exact_completion(mut self, required: bool) -> Self {
875        self.exact_completion = required;
876        self
877    }
878    /// Returns the requested topology, where `None` means replicated.
879    pub const fn topology(&self) -> Option<ParallelTopology> {
880        self.topology
881    }
882    /// Returns the requested weight residency.
883    pub const fn residency(&self) -> LayerWeightResidency {
884        self.residency
885    }
886    /// Returns the requested state policy.
887    pub const fn state(&self) -> &CacheResidencyPolicy {
888        &self.state
889    }
890    /// Returns the requested transform.
891    pub const fn quantization(&self) -> Option<QuantizationRequest> {
892        self.quantization
893    }
894    /// Returns requested session facilities.
895    pub const fn session(&self) -> SessionCapabilities {
896        self.session
897    }
898    /// Returns whether prompt-cache persistence is requested.
899    pub const fn prompt_cache(&self) -> bool {
900        self.prompt_cache
901    }
902    /// Returns whether exact completion ownership is requested.
903    pub const fn exact_completion(&self) -> bool {
904        self.exact_completion
905    }
906}
907
908/// Selected lowering for one canonical logical parameter.
909#[derive(Debug, Clone, Eq, PartialEq)]
910pub struct SelectedParameterRealization {
911    /// Canonical logical parameter identity.
912    name: String,
913    /// Physical outputs admitted as sources for this logical parameter.
914    sources: Vec<String>,
915    physical_sources: Vec<ReplicatedTextPhysicalSource>,
916    /// Admitted physical encoding.
917    source_encoding: SourceTensorEncoding,
918    /// Exact executable format used to construct the architecture module.
919    executable: LinearFormat,
920    /// Backend lowering selected for materialization.
921    lowering: WeightLoweringKind,
922}
923
924impl SelectedParameterRealization {
925    /// Returns the canonical logical identity.
926    pub fn name(&self) -> &str {
927        &self.name
928    }
929    /// Returns admitted physical source identities.
930    pub fn sources(&self) -> &[String] {
931        &self.sources
932    }
933    /// Returns the exact selected shard and multi-output provenance.
934    pub fn physical_sources(&self) -> &[ReplicatedTextPhysicalSource] {
935        &self.physical_sources
936    }
937    /// Returns the admitted source encoding.
938    pub const fn source_encoding(&self) -> &SourceTensorEncoding {
939        &self.source_encoding
940    }
941    /// Returns the selected executable format.
942    pub const fn executable(&self) -> LinearFormat {
943        self.executable
944    }
945    /// Returns the selected backend lowering kind.
946    pub const fn lowering(&self) -> WeightLoweringKind {
947        self.lowering
948    }
949}
950
951/// Authoritative realization selected before architecture or payload construction.
952#[derive(Debug, Clone, Eq, PartialEq)]
953pub struct SelectedReplicatedTextRealization {
954    /// Exact selected execution topology.
955    topology: ParallelTopology,
956    /// Selected ordinary parameter residency.
957    residency: LayerWeightResidency,
958    /// Selected mutable-state implementation and its exact residency policy.
959    state: CacheResidencyPolicy,
960    /// Exact per-parameter source, executable format, and lowering.
961    parameters: Vec<SelectedParameterRealization>,
962    /// Required observation facilities admitted by the backend.
963    session: SessionCapabilities,
964    /// Prompt-cache persistence is selected for this lifecycle.
965    prompt_cache: bool,
966    /// Exact completion ownership selected for this lifecycle.
967    exact_completion: bool,
968    grouped_operations: Vec<GroupedOperationRequirement>,
969}
970
971impl SelectedReplicatedTextRealization {
972    /// Returns the exact selected topology.
973    pub const fn topology(&self) -> ParallelTopology {
974        self.topology
975    }
976    /// Returns selected weight residency.
977    pub const fn residency(&self) -> LayerWeightResidency {
978        self.residency
979    }
980    /// Returns selected mutable-state policy.
981    pub const fn state(&self) -> &CacheResidencyPolicy {
982        &self.state
983    }
984    /// Returns exact per-parameter realizations.
985    pub fn parameters(&self) -> &[SelectedParameterRealization] {
986        &self.parameters
987    }
988    /// Returns selected session facilities.
989    pub const fn session(&self) -> SessionCapabilities {
990        self.session
991    }
992    /// Returns whether prompt-cache persistence was selected.
993    pub const fn prompt_cache(&self) -> bool {
994        self.prompt_cache
995    }
996    /// Returns whether exact completion ownership was selected.
997    pub const fn exact_completion(&self) -> bool {
998        self.exact_completion
999    }
1000    /// Returns selected grouped operation mechanisms.
1001    pub fn grouped_operations(&self) -> &[GroupedOperationRequirement] {
1002        &self.grouped_operations
1003    }
1004}
1005
1006/// Complete fail-closed selection diagnostic.
1007#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1008#[error("replicated text realization is unsupported: {issues}", issues = .issues.join("; "))]
1009pub struct ReplicatedTextSelectionError {
1010    issues: Vec<String>,
1011}
1012
1013impl ReplicatedTextSelectionError {
1014    /// Every missing semantic or mechanism requirement in stable order.
1015    pub fn issues(&self) -> &[String] {
1016        &self.issues
1017    }
1018}
1019
1020/// Deterministically selects one realization without constructing backend payloads.
1021pub fn select_replicated_text_realization(
1022    requirements: &ReplicatedTextRequirements,
1023    request: &ReplicatedTextSelectionRequest,
1024    capabilities: &BackendMechanismCapabilities,
1025) -> Result<SelectedReplicatedTextRealization, ReplicatedTextSelectionError> {
1026    let mut issues = Vec::new();
1027    if request
1028        .topology
1029        .is_some_and(|topology| !topology.is_replicated())
1030    {
1031        issues.push("replicated execution topology".into());
1032    }
1033    if !capabilities.operators.contains(requirements.operators) {
1034        issues.extend(
1035            capabilities
1036                .operators
1037                .missing_capability_names(requirements.operators)
1038                .into_iter()
1039                .map(|name| format!("neural operation {name}")),
1040        );
1041    }
1042    for operation in &requirements.grouped_operations {
1043        if !capabilities.grouped_operations.contains(operation) {
1044            issues.push(format!("grouped operation {operation:?}"));
1045        }
1046    }
1047    let residency_mechanism = match request.residency {
1048        LayerWeightResidency::FullyResident => WeightResidencyMechanism::Resident,
1049        LayerWeightResidency::LayerwiseHost(_) => WeightResidencyMechanism::Windowed,
1050        LayerWeightResidency::DenseDiskStream(_) => WeightResidencyMechanism::DiskStreamed,
1051    };
1052    if !capabilities
1053        .weight_residencies
1054        .contains(&residency_mechanism)
1055    {
1056        issues.push(format!("weight residency {residency_mechanism:?}"));
1057    }
1058    let state_residency = match &request.state {
1059        CacheResidencyPolicy::Device => StateResidencyMechanism::Device,
1060        CacheResidencyPolicy::Paged(_) => StateResidencyMechanism::Paged,
1061    };
1062    if !capabilities.state_residencies.contains(&state_residency) {
1063        issues.push(format!("state residency {state_residency:?}"));
1064    }
1065    for (required, supported, name) in [
1066        (
1067            request.session.persistent_cache(),
1068            capabilities.session.persistent_cache(),
1069            "persistent_cache",
1070        ),
1071        (
1072            request.session.output_observation(),
1073            capabilities.session.output_observation(),
1074            "output_observation",
1075        ),
1076        (
1077            request.session.activation_inspection(),
1078            capabilities.session.activation_inspection(),
1079            "activation_inspection",
1080        ),
1081    ] {
1082        if required && !supported {
1083            issues.push(format!("session capability {name}"));
1084        }
1085    }
1086    if request.prompt_cache && !capabilities.prompt_cache {
1087        issues.push("prompt-cache persistence".into());
1088    }
1089    if request.exact_completion && !capabilities.exact_completion {
1090        issues.push("exact completion ownership".into());
1091    }
1092
1093    let mut parameters = Vec::with_capacity(requirements.parameters.len());
1094    let mut names = BTreeSet::new();
1095    for parameter in &requirements.parameters {
1096        if parameter.name.trim().is_empty() || !names.insert(parameter.name.as_str()) {
1097            issues.push(format!(
1098                "unique nonempty logical parameter identity {:?}",
1099                parameter.name
1100            ));
1101            continue;
1102        }
1103        if !parameter.presence.has_physical_source() {
1104            continue;
1105        }
1106        let candidate = match request.quantization {
1107            Some(request) => match parameter.transform_target(request) {
1108                Ok(Some(target)) => Some((target.executable(), target.descriptor().clone())),
1109                Ok(None) => Some((
1110                    parameter.native_executable,
1111                    parameter
1112                        .lowering_descriptor(parameter.native_executable)
1113                        .expect("validated parameter forms native descriptor"),
1114                )),
1115                Err(error) => {
1116                    issues.push(error.to_string());
1117                    None
1118                }
1119            },
1120            None => Some((
1121                parameter.native_executable,
1122                parameter
1123                    .lowering_descriptor(parameter.native_executable)
1124                    .expect("validated parameter forms native descriptor"),
1125            )),
1126        };
1127        let Some((executable, descriptor)) = candidate else {
1128            issues.push(format!(
1129                "architecture transform {:?} for {:?}",
1130                request.quantization, parameter.name
1131            ));
1132            continue;
1133        };
1134        let Some(lowering) = capabilities
1135            .weight_lowerings
1136            .iter()
1137            .find(|lowering| lowering.descriptor == descriptor)
1138        else {
1139            issues.push(format!(
1140                "weight lowering {:?} -> {:?} for {:?}",
1141                parameter.source_encoding, executable, parameter.name
1142            ));
1143            continue;
1144        };
1145        parameters.push(SelectedParameterRealization {
1146            name: parameter.name.clone(),
1147            sources: parameter.sources.clone(),
1148            physical_sources: parameter.physical_sources.clone(),
1149            source_encoding: parameter
1150                .source_encoding
1151                .clone()
1152                .expect("physical parameter has a source encoding"),
1153            executable,
1154            lowering: lowering.kind,
1155        });
1156    }
1157    if !issues.is_empty() {
1158        return Err(ReplicatedTextSelectionError { issues });
1159    }
1160    Ok(SelectedReplicatedTextRealization {
1161        topology: request
1162            .topology
1163            .unwrap_or_else(|| ParallelTopology::new(1, 1, 1, 1).expect("replicated topology")),
1164        residency: request.residency,
1165        state: request.state.clone(),
1166        parameters,
1167        session: request.session,
1168        prompt_cache: request.prompt_cache,
1169        exact_completion: request.exact_completion,
1170        grouped_operations: requirements.grouped_operations.clone(),
1171    })
1172}
1173
1174#[cfg(test)]
1175mod tests {
1176    use super::*;
1177    use crate::{
1178        ArchitectureGroupKind, ArchitectureGroupPlacement, ArchitectureGroupTransport,
1179        ArchitectureMergeDestination, DenseDiskStreamLoadOptions, ExecutionGroupSpec,
1180        ExecutionUnitLayout, LayerwiseLoadOptions, StateLayout,
1181    };
1182    use eredu_checkpoint::{AffineQuantization, StoredDtype};
1183    use eredu_core::{cache::LayerCachePolicy, AttentionPolicy, LayerSchedule};
1184
1185    fn paged_state() -> CacheResidencyPolicy {
1186        CacheResidencyPolicy::Paged(
1187            crate::PagedCacheOptions::new(4, 1 << 20, 1 << 20, 1)
1188                .unwrap()
1189                .with_full_attention(true),
1190        )
1191    }
1192
1193    fn physical_source(name: &str) -> ReplicatedTextPhysicalSource {
1194        ReplicatedTextPhysicalSource::new(name, "/checkpoint/model.safetensors", name).unwrap()
1195    }
1196
1197    fn requirements() -> ReplicatedTextRequirements {
1198        let graph =
1199            ExecutionGraph::new(vec![ExecutionGroupSpec::root("decoder")], "decoder").unwrap();
1200        let execution_units = ExecutionUnitLayout::new(&graph, [1]).unwrap();
1201        ReplicatedTextRequirements::new(
1202            NeuralOperatorCapabilities::EXP,
1203            graph,
1204            execution_units,
1205            vec![ArchitectureGroupTransport {
1206                placement: ArchitectureGroupPlacement::Pipeline,
1207                kind: ArchitectureGroupKind::Decoder,
1208                first_owner_static_roles: vec!["embedding".into()],
1209                last_owner_static_roles: vec!["output".into()],
1210                merge_destination: ArchitectureMergeDestination::LastOwner,
1211                parallel_subgroup: None,
1212                request_optional: false,
1213            }],
1214            StateLayout::new(
1215                LayerSchedule::new(
1216                    1,
1217                    vec![LayerCachePolicy::key_value(AttentionPolicy::Full, 1, 8).unwrap()],
1218                )
1219                .unwrap(),
1220            )
1221            .unwrap(),
1222            vec![
1223                ReplicatedTextParameterRequirement::new(
1224                    "model.layers.0.mlp.weight",
1225                    vec!["blk.0.ffn.weight".into()],
1226                    vec![physical_source("blk.0.ffn.weight")],
1227                    Vec::new(),
1228                    Some(SourceTensorEncoding::Safetensors(StoredDtype::F16)),
1229                    Some(vec![64, 64]),
1230                    vec![64, 64],
1231                    LinearFormat::Dense,
1232                    ReplicatedTextParameterRole::LinearWeight,
1233                    ReplicatedTextParameterOwner::ExecutionUnit {
1234                        group: "decoder".into(),
1235                        unit: 0,
1236                    },
1237                    ReplicatedTextParameterPresence::Required,
1238                    ParameterTransformConstraint::Linear { packed_axis: 1 },
1239                )
1240                .unwrap(),
1241                ReplicatedTextParameterRequirement::new(
1242                    "model.layers.0.mlp.bias",
1243                    Vec::new(),
1244                    Vec::new(),
1245                    Vec::new(),
1246                    None,
1247                    None,
1248                    vec![64],
1249                    LinearFormat::Dense,
1250                    ReplicatedTextParameterRole::LinearBias,
1251                    ReplicatedTextParameterOwner::ExecutionUnit {
1252                        group: "decoder".into(),
1253                        unit: 0,
1254                    },
1255                    ReplicatedTextParameterPresence::OptionalAbsent,
1256                    ParameterTransformConstraint::None,
1257                )
1258                .unwrap(),
1259                ReplicatedTextParameterRequirement::new(
1260                    "model.layers.0.norm.weight",
1261                    vec!["blk.0.norm.weight".into()],
1262                    vec![physical_source("blk.0.norm.weight")],
1263                    Vec::new(),
1264                    Some(SourceTensorEncoding::Safetensors(StoredDtype::F16)),
1265                    Some(vec![64]),
1266                    vec![64],
1267                    LinearFormat::Dense,
1268                    ReplicatedTextParameterRole::Normalization,
1269                    ReplicatedTextParameterOwner::ExecutionUnit {
1270                        group: "decoder".into(),
1271                        unit: 0,
1272                    },
1273                    ReplicatedTextParameterPresence::Required,
1274                    ParameterTransformConstraint::None,
1275                )
1276                .unwrap(),
1277            ],
1278        )
1279        .unwrap()
1280    }
1281
1282    #[test]
1283    fn parameter_requirement_preserves_every_admitted_alias() {
1284        let requirement = ReplicatedTextParameterRequirement::new(
1285            "model.layers.0.mlp.weight",
1286            vec!["released.layers.0.mlp.weight".into()],
1287            vec![physical_source("released.layers.0.mlp.weight")],
1288            vec![
1289                "legacy.layers.0.mlp.weight".into(),
1290                "vendor.layers.0.mlp.weight".into(),
1291            ],
1292            Some(SourceTensorEncoding::Safetensors(StoredDtype::F16)),
1293            Some(vec![64, 64]),
1294            vec![64, 64],
1295            LinearFormat::Dense,
1296            ReplicatedTextParameterRole::LinearWeight,
1297            ReplicatedTextParameterOwner::ExecutionUnit {
1298                group: "decoder".into(),
1299                unit: 0,
1300            },
1301            ReplicatedTextParameterPresence::Required,
1302            ParameterTransformConstraint::Linear { packed_axis: 1 },
1303        )
1304        .unwrap();
1305
1306        assert_eq!(
1307            requirement.aliases(),
1308            ["legacy.layers.0.mlp.weight", "vendor.layers.0.mlp.weight"]
1309        );
1310        assert_eq!(requirement.sources(), ["released.layers.0.mlp.weight"]);
1311
1312        let absent_bias = ReplicatedTextParameterRequirement::new(
1313            "model.layers.0.mlp.bias",
1314            Vec::new(),
1315            Vec::new(),
1316            vec!["released.layers.0.mlp.bias".into()],
1317            None,
1318            None,
1319            vec![64],
1320            LinearFormat::Dense,
1321            ReplicatedTextParameterRole::LinearBias,
1322            ReplicatedTextParameterOwner::ExecutionUnit {
1323                group: "decoder".into(),
1324                unit: 0,
1325            },
1326            ReplicatedTextParameterPresence::OptionalAbsent,
1327            ParameterTransformConstraint::None,
1328        )
1329        .unwrap();
1330        assert_eq!(
1331            absent_bias.presence(),
1332            &ReplicatedTextParameterPresence::OptionalAbsent
1333        );
1334        assert!(absent_bias.sources().is_empty());
1335        assert_eq!(
1336            absent_bias.transform_constraint(),
1337            ParameterTransformConstraint::None
1338        );
1339    }
1340
1341    #[test]
1342    fn physical_provenance_distinguishes_outputs_from_one_sharded_tensor() {
1343        let shard = "/checkpoint/model-00002-of-00003.gguf";
1344        let weight = ReplicatedTextPhysicalSource::new(
1345            "blk.0.ffn_gate.weight",
1346            shard,
1347            "blk.0.ffn_gate.weight",
1348        )
1349        .unwrap();
1350        let scales = ReplicatedTextPhysicalSource::new(
1351            "blk.0.ffn_gate.weight",
1352            shard,
1353            "blk.0.ffn_gate.scales",
1354        )
1355        .unwrap();
1356        assert_eq!(weight.tensor(), scales.tensor());
1357        assert_eq!(weight.shard(), scales.shard());
1358        assert_ne!(weight.output(), scales.output());
1359    }
1360
1361    fn capabilities() -> BackendMechanismCapabilities {
1362        let source = SourceTensorEncoding::Safetensors(StoredDtype::F16);
1363        BackendMechanismCapabilities::new(
1364            NeuralOperatorCapabilities::EXP,
1365            vec![
1366                WeightLoweringCapability::new(
1367                    WeightLoweringDescriptor::new(
1368                        source.clone(),
1369                        LinearFormat::Dense,
1370                        vec![64, 64],
1371                        vec![64, 64],
1372                        Some(1),
1373                    )
1374                    .unwrap(),
1375                    WeightLoweringKind::Direct,
1376                ),
1377                WeightLoweringCapability::new(
1378                    WeightLoweringDescriptor::new(
1379                        source,
1380                        LinearFormat::Affine(AffineQuantization::new(64, 4).unwrap()),
1381                        vec![64, 64],
1382                        vec![64, 64],
1383                        Some(1),
1384                    )
1385                    .unwrap(),
1386                    WeightLoweringKind::Transform,
1387                ),
1388                WeightLoweringCapability::new(
1389                    WeightLoweringDescriptor::new(
1390                        SourceTensorEncoding::Safetensors(StoredDtype::F16),
1391                        LinearFormat::Dense,
1392                        vec![64],
1393                        vec![64],
1394                        None,
1395                    )
1396                    .unwrap(),
1397                    WeightLoweringKind::Direct,
1398                ),
1399            ],
1400            vec![
1401                WeightResidencyMechanism::Resident,
1402                WeightResidencyMechanism::Windowed,
1403                WeightResidencyMechanism::DiskStreamed,
1404            ],
1405            vec![
1406                StateResidencyMechanism::Device,
1407                StateResidencyMechanism::Paged,
1408            ],
1409        )
1410        .with_session(SessionCapabilities::new(true, true, true))
1411        .with_prompt_cache(true)
1412        .with_exact_completion(true)
1413    }
1414
1415    fn request(residency: LayerWeightResidency) -> ReplicatedTextSelectionRequest {
1416        ReplicatedTextSelectionRequest::new(residency, paged_state())
1417            .with_session(SessionCapabilities::new(true, true, true))
1418            .with_prompt_cache(true)
1419            .with_exact_completion(true)
1420    }
1421
1422    #[test]
1423    fn complete_requirements_are_invariant_across_all_caller_policy_dimensions() {
1424        let baseline = requirements();
1425        let disk = DenseDiskStreamLoadOptions::new(4096, 8192, 2, 1).unwrap();
1426        let requests = [
1427            ReplicatedTextSelectionRequest::new(
1428                LayerWeightResidency::FullyResident,
1429                CacheResidencyPolicy::Device,
1430            ),
1431            ReplicatedTextSelectionRequest::new(
1432                LayerWeightResidency::LayerwiseHost(LayerwiseLoadOptions::default()),
1433                paged_state(),
1434            )
1435            .with_topology(ParallelTopology::new(2, 1, 1, 1).unwrap())
1436            .with_quantization(QuantizationRequest::Affine {
1437                group_size: 64,
1438                bits: 4,
1439            })
1440            .with_session(SessionCapabilities::new(true, true, true))
1441            .with_prompt_cache(true)
1442            .with_exact_completion(true),
1443            ReplicatedTextSelectionRequest::new(
1444                LayerWeightResidency::DenseDiskStream(disk),
1445                CacheResidencyPolicy::Device,
1446            )
1447            .with_quantization(QuantizationRequest::MxFp4),
1448        ];
1449
1450        for _request in &requests {
1451            assert_eq!(requirements(), baseline);
1452        }
1453        assert_eq!(requests[0].state(), &CacheResidencyPolicy::Device);
1454        assert!(matches!(
1455            requests[1].residency(),
1456            LayerWeightResidency::LayerwiseHost(_)
1457        ));
1458        assert_eq!(requests[1].topology().unwrap().tensor(), 2);
1459        assert_eq!(
1460            requests[1].quantization(),
1461            Some(QuantizationRequest::Affine {
1462                group_size: 64,
1463                bits: 4,
1464            })
1465        );
1466        assert!(requests[1].prompt_cache());
1467        assert!(requests[1].exact_completion());
1468        assert!(requests[1].session().activation_inspection());
1469        assert_eq!(
1470            requests[2].residency(),
1471            LayerWeightResidency::DenseDiskStream(disk)
1472        );
1473        assert_eq!(requests[2].quantization(), Some(QuantizationRequest::MxFp4));
1474    }
1475
1476    #[test]
1477    fn selection_is_deterministic_and_keeps_source_format_distinct() {
1478        let disk = DenseDiskStreamLoadOptions::new(1234, 5678, 3, 2).unwrap();
1479        let request = request(LayerWeightResidency::DenseDiskStream(disk)).with_quantization(
1480            QuantizationRequest::Affine {
1481                group_size: 64,
1482                bits: 4,
1483            },
1484        );
1485        let left =
1486            select_replicated_text_realization(&requirements(), &request, &capabilities()).unwrap();
1487        let right =
1488            select_replicated_text_realization(&requirements(), &request, &capabilities()).unwrap();
1489        assert_eq!(left, right);
1490        assert_eq!(
1491            left.residency(),
1492            LayerWeightResidency::DenseDiskStream(disk)
1493        );
1494        assert_eq!(left.state(), &paged_state());
1495        assert_eq!(left.parameters().len(), 2);
1496        assert_eq!(requirements().parameters().len(), 3);
1497        assert!(matches!(
1498            requirements().parameters()[1].presence(),
1499            ReplicatedTextParameterPresence::OptionalAbsent
1500        ));
1501        assert!(matches!(
1502            requirements().parameters()[2].role(),
1503            ReplicatedTextParameterRole::Normalization
1504        ));
1505        assert_eq!(requirements().parameters()[2].logical_shape(), [64]);
1506        assert_eq!(
1507            requirements().parameters()[2].transform_constraint(),
1508            ParameterTransformConstraint::None
1509        );
1510        assert_eq!(
1511            left.parameters()[0].lowering(),
1512            WeightLoweringKind::Transform
1513        );
1514        assert_ne!(
1515            format!("{:?}", left.parameters()[0].source_encoding()),
1516            format!("{:?}", left.parameters()[0].executable())
1517        );
1518    }
1519
1520    #[test]
1521    fn selection_reports_all_missing_mechanisms_together() {
1522        let capabilities = BackendMechanismCapabilities::new(
1523            NeuralOperatorCapabilities::NONE,
1524            Vec::new(),
1525            Vec::new(),
1526            Vec::new(),
1527        );
1528        let error = select_replicated_text_realization(
1529            &requirements(),
1530            &request(LayerWeightResidency::LayerwiseHost(
1531                LayerwiseLoadOptions::default(),
1532            )),
1533            &capabilities,
1534        )
1535        .unwrap_err();
1536        assert!(error.issues().len() >= 7, "{:?}", error.issues());
1537        assert!(error.issues().iter().any(|issue| issue.contains("exp")));
1538        assert!(error
1539            .issues()
1540            .iter()
1541            .any(|issue| issue.contains("weight lowering")));
1542    }
1543
1544    #[test]
1545    fn transform_selection_rejects_incompatible_exact_geometry() {
1546        for quantization in [
1547            QuantizationRequest::Affine {
1548                group_size: 96,
1549                bits: 4,
1550            },
1551            QuantizationRequest::Affine {
1552                group_size: 256,
1553                bits: 4,
1554            },
1555            QuantizationRequest::Affine {
1556                group_size: 0,
1557                bits: 4,
1558            },
1559            QuantizationRequest::Affine {
1560                group_size: u32::MAX,
1561                bits: 4,
1562            },
1563            QuantizationRequest::Affine {
1564                group_size: 32,
1565                bits: 0,
1566            },
1567            QuantizationRequest::Affine {
1568                group_size: 32,
1569                bits: 7,
1570            },
1571        ] {
1572            let error = select_replicated_text_realization(
1573                &requirements(),
1574                &request(LayerWeightResidency::FullyResident).with_quantization(quantization),
1575                &capabilities(),
1576            )
1577            .unwrap_err();
1578            assert!(error
1579                .issues()
1580                .iter()
1581                .any(|issue| issue.contains("invalid replicated text contract")));
1582        }
1583
1584        let mut indivisible = requirements();
1585        indivisible.parameters[0].logical_shape = vec![64, 48];
1586        let error = select_replicated_text_realization(
1587            &indivisible,
1588            &request(LayerWeightResidency::FullyResident)
1589                .with_quantization(QuantizationRequest::MxFp4),
1590            &capabilities(),
1591        )
1592        .unwrap_err();
1593        assert!(error
1594            .issues()
1595            .iter()
1596            .any(|issue| issue.contains("MXFP4 packed extent 48")));
1597    }
1598
1599    #[test]
1600    fn exact_source_and_physical_geometry_fail_before_construction_or_payload() {
1601        for mutate in [
1602            |requirement: &mut ReplicatedTextParameterRequirement| {
1603                requirement.source_encoding =
1604                    Some(SourceTensorEncoding::Safetensors(StoredDtype::U8));
1605            },
1606            |requirement: &mut ReplicatedTextParameterRequirement| {
1607                requirement.physical_shape = Some(vec![64, 32]);
1608            },
1609        ] {
1610            let mut requirements = requirements();
1611            mutate(&mut requirements.parameters[0]);
1612            let selected = select_replicated_text_realization(
1613                &requirements,
1614                &request(LayerWeightResidency::FullyResident),
1615                &capabilities(),
1616            );
1617            let error = selected.unwrap_err();
1618            assert!(error
1619                .issues()
1620                .iter()
1621                .any(|issue| issue.contains("weight lowering")));
1622        }
1623    }
1624
1625    #[test]
1626    fn missing_tensor_parallel_grouped_partial_fails_before_construction_or_forward() {
1627        let requirements = requirements().with_grouped_operations([
1628            GroupedOperationRequirement::GatedProduct,
1629            GroupedOperationRequirement::GatedProductTensorParallelPartial,
1630        ]);
1631        let capabilities =
1632            capabilities().with_grouped_operations([GroupedOperationRequirement::GatedProduct]);
1633        let selected = select_replicated_text_realization(
1634            &requirements,
1635            &request(LayerWeightResidency::FullyResident),
1636            &capabilities,
1637        );
1638        let error = selected.unwrap_err();
1639        assert!(error
1640            .issues()
1641            .iter()
1642            .any(|issue| { issue.contains("GatedProductTensorParallelPartial") }));
1643    }
1644}
eredu_runtime/replicated_text.rs

eredu_runtime/
replicated_text.rs
