1#![warn(missing_docs)]
7
8pub mod automatic_support;
10pub mod backend;
12pub mod cache;
14pub mod capture;
16pub mod communication;
18pub mod component;
20pub mod composite;
22pub mod decision;
24pub mod dense;
26pub mod draft;
28pub mod execution;
30pub mod execution_control;
32pub mod execution_plan;
33pub mod expert;
34pub mod generation;
36pub mod input;
38pub mod inspection;
39pub mod intervention;
41pub mod layered;
43pub mod load_request;
45pub mod mechanism_synthesis;
47pub mod parallel;
49pub mod parameter;
51pub mod partition;
53pub mod partitioned_execution;
55pub mod placement;
56pub mod prefetch;
58pub mod realtime;
60pub mod realtime_executor;
62pub mod realtime_ingress;
64pub mod realtime_interpreter;
66pub mod realtime_mechanism_synthesis;
67pub mod realtime_model;
69pub mod realtime_payload;
71pub mod realtime_payload_state;
73pub mod realtime_selection;
75pub mod realtime_session;
77pub mod replicated_session;
79pub mod replicated_text;
81pub mod residency;
83pub mod speculative;
85pub mod speculative_selection;
87pub mod state;
89mod weight_residency;
90
91pub use automatic_support::{
92 residency_telemetry, selected_text_bounded_requirement, BoundedResidencySizingError,
93};
94pub use backend::{
95 BarrierBackend, BroadcastBackend, CollectiveBackend, CommunicationBackend, EvenGatherBackend,
96 FailureAgreementBackend, ParameterBackend, PointToPointBackend, RoleExactBoundaryValue,
97 SubmissionBackend, SumReductionBackend, TransferBackend, UnevenGatherBackend,
98 VariableAllToAllBackend,
99};
100pub use cache::{
101 finalize_prompt_cache_shard, hash_prompt_cache_shard_payload, inspect_prompt_cache,
102 prompt_cache_rank_path, resolve_prompt_cache_root, safe_prompt_cache_shard_path,
103 validate_prompt_cache_manifest, CacheBlockLifecycle, CacheBlockStorage,
104 CacheHostDemotionOperation, CacheHostPromotion, CacheIoAdmission, CacheIoCompletionDisposition,
105 CacheIoExecutionState, CacheIoExecutionStateError, CacheIoOperation, CacheIoOperationKey,
106 CacheIoOperationKind, CacheIoPreparation, CacheIoStartDisposition, CacheIoSubmission,
107 CacheIoSubmissionOutcome, CacheIoTicket, CacheIoWorker, CacheIoWorkerError,
108 CacheLayerResidencyReport, CacheLayerResidencyStats, CacheLifecycleError, CachePoolError,
109 CachePoolLimits, CachePoolMembership, CachePoolReport, CachePoolReservation, CachePoolResource,
110 CachePoolUsage, CacheResidencyConfigurationError, CacheResidencyPolicy, CacheResidencyPool,
111 CacheResidencyReport, CacheResidencyTelemetry, CacheStorageError, CacheStoragePhase,
112 LiveCacheBlockPublication, LiveCacheDiskPolicy, LiveCachePublicationError, MutableCacheTail,
113 PagedCacheOptions, PromptCachePersistenceError, PromptCachePublication,
114 ReversiblePromptCachePublication, CACHE_RESIDENCY_LAYER_REPORT_LIMIT,
115 MAX_PROMPT_CACHE_SHARD_HEADER_BYTES, PROMPT_CACHE_CURRENT_FILE,
116 PROMPT_CACHE_GENERATIONS_DIRECTORY,
117};
118pub use communication::validate_communication_manifest_consensus;
119pub use communication::{
120 prepare_communication_realization, project_all_communication_manifests,
121 project_communication_manifest, validate_compatible_communication_manifests,
122 BoundaryDimensionContract, BoundaryFramingProtocol, BoundaryRoleContract,
123 CommunicationCapabilities, CommunicationCapabilityError, CommunicationCompletionCapabilities,
124 CommunicationCompletionPolicy, CommunicationGroupDescriptor, CommunicationGroupRequirements,
125 CommunicationManifest, CommunicationManifestConsensusError, CommunicationManifestError,
126 CommunicationOperation, CommunicationOperationRequirement, CommunicationPeerCounts,
127 CommunicationRealizationError, CommunicationRouteDescriptor, CommunicationRouteId,
128 CommunicationTensorLimits, CommunicationTopologyCapabilities, PreparedCommunicationRealization,
129 RoleExactBoundaryContract, TopologyCommunicationPlan,
130};
131pub use component::{
132 ComponentDomain, ComponentGraph, ComponentGraphError, ComponentKind, ComponentResidencyClass,
133 ComponentSpec,
134};
135pub use composite::{
136 select_composite_realization, select_processor_execution, CompositeSelectionError,
137 MediaPrimitiveCapabilities, ModalityProcessorRequirements, ProcessorExecutionRequirements,
138 ProcessorPrimitive, ProcessorSelectionError, ProcessorSelectionRequest,
139 SelectedCompositeRealization, SelectedProcessorExecution,
140};
141pub use decision::{
142 FullyForcedTailDecision, PredictionDirective, SequentialDecision, SequentialDecisionBoundary,
143 SequentialDecisionDiagnostic, SequentialDecisionDriver, SequentialDecisionError,
144 SequentialDecisionMode, SequentialDecisionPlan, SequentialDecisionPlanError,
145 SequentialDecisionSource, SequentialDecisionTraversal, SequentialSamplingState,
146};
147pub use dense::{
148 DenseCacheMetrics, DenseDiskStreamReport, DenseExecutionGroupReport, DensePassCounterSnapshot,
149 DensePassReport, DenseStreamTelemetry, DenseStreamTelemetryError, DenseTierResidencyReport,
150};
151pub use draft::{execute_draft_group, DraftGroupExecutionError, DraftStateTransaction};
152pub use execution::{
153 ExecutionGraph, ExecutionGraphError, ExecutionGroupId, ExecutionGroupSchedule,
154 ExecutionGroupSpec, ExecutionScheduleError, ExecutionUnitAddress, ExecutionUnitLayout,
155 ExecutionUnitLayoutError, ReadyGroupState,
156};
157pub use execution_plan::{
158 execution_plan_quantization, ExecutionPlanLoadError, ResidencyDiagnostics,
159};
160pub use expert::{
161 combine_routed_expert_tensor_parallel, combine_tensor_parallel_expert_outputs,
162 plan_addressable_bank_bindings, reduce_routed_expert_tensor_parallel,
163 reduce_tensor_parallel_expert_output, selected_addressable_parameter_bytes,
164 AddressableBankBindingPlan, AddressableBankDistribution, AddressableBankMember,
165 AddressableBankMemberError, AddressableBankMemberPlacement, AddressableBankParameter,
166 AddressableBindingTransform, AddressableExpertRouteProvider, AddressableExpertRouteRequest,
167 AddressableGatedProductBank, AddressableGroupedBank, ExpertRouteCombination,
168 ExpertRouteExchange, ExpertRouteTensorMovement, IndexedMovement, ObservedExpertProvider,
169 ObservedExpertProviderError, ParameterBankAcquisition, ResidentExpertProvider,
170 RoutedExpertProvider, RoutedExpertRequest, RoutedExpertTensorParallelOutput,
171 RoutedObservationPoint, TensorParallelRoutedExpertProvider,
172};
173pub use expert::{select_routes_with_observer, select_routes_with_provider};
174pub use generation::{
175 CausalModel, ConstrainedSampler, DefaultSampler, GenerationSampler, MirostatV2Sampler,
176 PenaltyConfig, Sampler, SamplingBackend, SamplingConfigurationError, SpeculativeSampler,
177 TokenDomain,
178};
179pub use input::{
180 PreparedInputCacheIdentity, PreparedInputCacheIdentityError, PreparedInputInspector,
181 PreparedInputPart, PreparedInputPayload, PreparedModelInput,
182};
183pub use inspection::{
184 observe_and_intervene, observe_model_logits, ActivationObserver, NoopObserver, RoutingDecision,
185 RoutingObservation, TargetStateCapture, TargetStateCaptureError, TargetStateTap,
186};
187pub use layered::{
188 ArchitectureGroupKind, ArchitectureGroupPlacement, ArchitectureGroupTransport,
189 ArchitectureMergeDestination, ArchitectureParallelSubgroup, ArchitectureParameters,
190 CompositeLayeredTraversalHook, LayeredArchitecture, LayeredForwardState, LayeredPartitionInput,
191 LayeredPartitionOutput, LayeredPipelineSchedule, LayeredPipelineScheduleError,
192 LayeredTraversalHook, LayeredTraversalPoint, LayeredUnitAction, LayerwiseAcquireError,
193 LayerwisePolicy, LayerwisePolicyForward, LayerwiseRuntime, LayerwiseRuntimeError,
194 ParallelLayeredArchitecture, ParallelRoutedLayeredArchitecture, PartitionedLayeredArchitecture,
195 ResidentRuntime, ResidentUnitWindow, ResidentUnitWindowError, RoutedLayeredArchitecture,
196 StaticParameterVisitor, StaticParameterVisitorMut,
197};
198pub use load_request::{
199 DraftingLoadRequest, NormalizedLoadRequest, NormalizedLoadRequestError, ParallelLoadRequest,
200 ValidatedModelLoadRequest,
201};
202pub use mechanism_synthesis::{
203 synthesize_replicated_text_capabilities, BackendMechanismFacts, ReplicatedTextMechanismSupport,
204 StateLifecycleCapabilities,
205};
206pub use parallel::{
207 aligned_partition_units, expand_linear_format_parameter_groups, module_parameter_group,
208 partitioned_module_parameter_group, partitioned_projection_group, projection_parameter_group,
209 segmented_projection_group, LocalModelLayout, LocalTensorLayout, MemberSharding,
210 ParallelModelInfo, ParallelPlanError, ParameterGroupSpec, ParameterMemberSpec, ParameterRole,
211 ProjectionSharding, ShardingPolicy, TensorPlacement,
212};
213pub use parameter::{
214 bind_materialized_unit, bind_materialized_unit_excluding, bindings_from_recipe_set,
215 build_exact_replicated_text_bindings, build_module_binding_plan, materialize_bindings,
216 materialize_selected_bindings, preflight_bindings, select_bindings, BindingPlan,
217 BindingPlanError, MaterializedUnit, ModuleBindingPlan, ModuleBindingPlanError,
218 ParameterBindingTarget, ParameterOrchestrationError, PlannedBinding, RecipeBindingError,
219 SelectedBindingPlan,
220};
221pub use partition::{
222 validate_boundary_tensor_count, ArchitectureBoundary, ArchitectureBoundaryError,
223 ArchitectureBoundaryValue, ArchitectureParameterDescription, ArchitectureParameterError,
224 ArchitecturePartition, ArchitecturePartitionError, BoundaryTensorDimension,
225 BoundaryTensorDtype, BoundaryTensorSpec, BoundaryWireSchema, LayeredPartitionBeginError,
226 LayeredPartitionDriver, LayeredPartitionError, NoAuxiliaryBoundary, NoAuxiliaryBoundarySchema,
227 OwnedParameterGroupSpec, ParameterGroupOwner, PartitionGroup, PartitionOwnership,
228 PartitionState, PipelineActivationDtype, PipelineWireContract, ResolvedBoundaryTensorSpec,
229 ResolvedBoundaryWireSchema,
230};
231pub use partitioned_execution::{
232 CommunicationTensorMetadata, DistributedExecutionPhase, LayerwiseTraversalPartitionExecutor,
233 LayerwiseTraversalRuntime, NoBoundaryTransport, NoCommitAgreement, NoOutputPublisher,
234 OpaqueBoundaryTransport, OpaqueCommitAgreement, OpaqueFailureAgreement, OpaqueOutputPublisher,
235 PartitionBoundaryRoute, PartitionBoundaryTransport, PartitionCommitAgreement,
236 PartitionCommunication, PartitionCommunicationAuthority, PartitionExecutionError,
237 PartitionOutputAuthority, PartitionOutputPublication, PartitionOutputPublisher,
238 PartitionedExecutionPlan, PartitionedGroupExecutor, PartitionedTextExecution,
239 PartitionedTextRuntime, PartitionedTraversalError, PartitionedTraversalResult,
240 RealizedCommunicationGroup, RealizedCommunicationRoute,
241};
242pub use placement::{
243 place_addressable_member_bindings, place_weight_bindings, placement_selection,
244 BindingPlacementError, PlacementPlan, PlacementPlanError, PlacementRank,
245 ResolvedTensorPlacement, TensorSlice,
246};
247pub use prefetch::{BackgroundPrefetchWorker, BackgroundPrefetchWorkerError};
248pub use realtime::{
249 RealtimeCompletionAttachmentError, RealtimeFrameExecutionError, RealtimeFrameTransition,
250 RealtimeGenerationBranch, RealtimeGenerationState, RealtimeGenerationTransactionError,
251};
252pub use realtime_executor::{
253 execute_realtime_frame, PreparedRealtimeFrameExecutor, PrepublicationRealtimeFrame,
254 RealtimeCompletionCreationError, RealtimeDecisionExecution, RealtimeFrameCompletionMechanism,
255 RealtimeFrameCoordinatorError, RealtimeFrameHostObserver, RealtimeHostOutputUnavailable,
256 RealtimePrepublicationError, SubmittedRealtimeFrame,
257};
258pub use realtime_ingress::{
259 MaterializedRealtimeInput, RealtimeHostTokenMaterializer, RealtimeIngressContract,
260 RealtimeIngressError, RealtimePayloadKind, RealtimeTokenKind, ValidatedRealtimeInput,
261};
262pub use realtime_interpreter::{
263 complete_realtime_frame, prepare_realtime_frame, CompletedRealtimeFrame, PreparedRealtimeFrame,
264 RealtimeFrameInterpretationError, RealtimeFrameTensorMechanisms,
265};
266pub use realtime_mechanism_synthesis::{
267 synthesize_realtime_capabilities, RealtimeMechanismFacts, RealtimeMechanismSupport,
268};
269pub use realtime_model::{
270 construct_realtime_model, preflight_realtime_materialization_tasks, realtime_task_binding_plan,
271 ConstructedRealtimeExecution, ConstructedRealtimeModel, PreparedRealtimeModelContract,
272 RealizedRealtimePolicy, RealizedRealtimeState, RealtimeArchitectureConstructionIdentity,
273 RealtimeArchitectureIdentity, RealtimeLayerwiseRuntime, RealtimeMaterializationComponent,
274 RealtimeMaterializationTask, RealtimeModelConstructionError,
275 RealtimeModelConstructionMechanisms, RealtimeModelContractError, RealtimeTaskBindingPlan,
276};
277pub use realtime_payload::{
278 RealtimePayloadContract, RealtimePayloadContractError, RealtimePayloadEnvelope,
279 RealtimePayloadGeneration, RealtimePayloadHistory, RealtimePayloadHistoryError,
280 RealtimePayloadOwnerIdentity,
281};
282pub use realtime_payload_state::{
283 RealtimePayloadBranch, RealtimePayloadState, RealtimePayloadStateTransactionError,
284};
285pub use realtime_selection::{
286 select_and_prepare_realtime_realization, select_realtime_realization,
287 ConstructedRealtimeResources, PreparedRealtimeRealization, RealtimeArchitectureProof,
288 RealtimeArchitectureRequirements, RealtimeContractError, RealtimeExecutionRequirements,
289 RealtimeIdentity, RealtimeMechanism, RealtimeMechanismCapabilities,
290 RealtimeMechanismRequirements, RealtimeObservationRequirements, RealtimePreparationError,
291 RealtimeSelectionError, RealtimeSelectionIssue, RealtimeSelectionRequest,
292 RealtimeTopologyPolicy, RealtimeWeightComponentRequirement, RealtimeWeightComponentRole,
293 RealtimeWeightLoweringRequirement, SelectedRealtimeRealization,
294 SelectedRealtimeStateComponentRealization, SelectedRealtimeStateRealization,
295};
296pub use realtime_session::{
297 RealtimeHistoryGeneration, RealtimeModelOwnerIdentity, RealtimeModelSessionIdentity,
298 RealtimeSamplingReplacementError, RealtimeSamplingUpdateError, RealtimeSessionBranch,
299 RealtimeSessionError, RealtimeSessionExecutionError, RealtimeSessionIncarnation,
300 RealtimeSessionResumeError, RealtimeSessionScheduler, RealtimeSessionState,
301 RealtimeSessionTransactionError, ReleasedRealtimeSession,
302};
303pub use replicated_session::{
304 construct_replicated_text_session, construct_replicated_text_session_with_execution,
305 construct_replicated_text_session_with_runtime, prepare_default_partitioned_runtime,
306 prepare_layered_text_contract, prepare_layered_text_contract_with_addressable_parameters,
307 prepare_partitioned_session_runtime, prepare_partitioned_session_runtime_with_exclusions,
308 prepare_replicated_text_contract, prepare_replicated_text_contract_with_addressable_parameters,
309 DirectReplicatedTextExecution, DistributedSessionCheckpoint, DistributedStateCheckpoint,
310 PartitionedRuntimeConstructionError, PartitionedSessionFactoryInput,
311 PartitionedSessionPreparationError, PartitionedUnitScope, PredictionTargetOperation,
312 PreparedPartitionedRuntimeComponents, PreparedPartitionedSessionRuntime,
313 PreparedReplicatedTextContract, ReplicatedRuntimeExecutionStrategy,
314 ReplicatedTextExecutionStrategy, ReplicatedTextRuntime, ReplicatedTextSession,
315 ReplicatedTextSessionCheckpoint, ReplicatedTextSessionError, ReplicatedTextSessionMechanisms,
316 ReplicatedTextSessionReport, RoutedReplicatedTextExecution, SessionStateRealization,
317 TransactionalPromptCacheMechanisms,
318};
319pub use replicated_text::{
320 group_replicated_text_transform_tasks, locally_materialized_replicated_text_outputs,
321 partition_selected_replicated_text_materialization_tasks,
322 partitioned_replicated_text_materialization_tasks,
323 plan_local_replicated_text_materialization_tasks, plan_replicated_text_materialization_tasks,
324 replicated_text_materialization_tasks, select_replicated_text_realization,
325 selected_materialization_task_bytes, AddressableStorageCapabilities, AddressableStorageTiers,
326 BackendMechanismCapabilities, GroupedOperationRequirement, ParameterTransformConstraint,
327 ParameterTransformTarget, ReplicatedTextArchitecture, ReplicatedTextContractError,
328 ReplicatedTextMaterializationPartitionPlan, ReplicatedTextMaterializationTask,
329 ReplicatedTextOutputCompanion, ReplicatedTextOutputSelection, ReplicatedTextParameterOwner,
330 ReplicatedTextParameterPresence, ReplicatedTextParameterRequirement,
331 ReplicatedTextParameterRole, ReplicatedTextPhysicalSource, ReplicatedTextRequirements,
332 ReplicatedTextSelectionError, ReplicatedTextSelectionRequest, ReplicatedTextStateAccess,
333 ReplicatedTextTransformGroup, SelectedParameterRealization, SelectedReplicatedTextRealization,
334 SelectedStateComponentRealization, SelectedStateRealization, StateComponentMechanism,
335 StateComponentPlacement, StateMechanismCapabilities, StateStorageDtype,
336 WeightLoweringCapability, WeightLoweringDescriptor, WeightLoweringKind,
337 WeightResidencyMechanism,
338};
339pub use residency::{
340 DeviceLayerWindow, OffloadUnit, QuantizationCompanionBindings, ResidencyAcquisition,
341 ResidencyController, ResidencyControllerError, ResidencyDeclarationError, ResidencyLease,
342 ResidencyLeaseOwner, ResidencyLeaseStorage, ResidencyReport, ResidencyTransfer,
343 ResidencyTransferOwner, ResidencyWindowError, ResidencyWindowManager, ResidentLayerGroup,
344 ResidentLayerGroupReport, WeightBinding, WeightBindingPlan, WeightBindingSelectionError,
345 WeightMaterializationReport,
346};
347pub use speculative::{RunSpeculativeGeneration, SpeculativeScheduler};
348pub use speculative_selection::{
349 select_and_prepare_speculative_realization,
350 select_and_prepare_speculative_realization_observed, select_speculative_realization,
351 ConstructedSpeculativeResources, PreparedSpeculativeRealization, SelectedSpeculativeCompletion,
352 SelectedSpeculativePlacement, SelectedSpeculativeRealization, SelectedSpeculativeSampling,
353 SelectedSpeculativeState, SpeculativeArchitectureCompatibilityProof, SpeculativeCaptureEntry,
354 SpeculativeCaptureEnvelope, SpeculativeCaptureError, SpeculativeCaptureMetadata,
355 SpeculativeCaptureSchema, SpeculativeContractError, SpeculativeIdentity,
356 SpeculativeLaneIdentity, SpeculativeMechanism, SpeculativeMechanismCapabilities,
357 SpeculativeMechanismRequirements, SpeculativePlacementRequest, SpeculativePreparationError,
358 SpeculativeRealizationRequirements, SpeculativeSelectionError, SpeculativeSelectionRequest,
359 SpeculativeStateCacheIdentityIngredients, SpeculativeStrategyClass,
360 SpeculativeStrategyRequirements,
361};
362pub use state::{
363 realize_architecture_state, ArchitectureStateFactory, ArchitectureStatePartitionError,
364 ArchitectureStatePartitionPlan, ArchitectureStatePartitionRule, ArchitectureStatePlacement,
365 ArchitectureStateRealizationError, DeviceState, LayerRuntimeState, ModelStateIdentity,
366 ResettableRuntimeLayerState, ResettableRuntimeState, RuntimeLayerState, RuntimeState,
367 RuntimeStateComponents, StateError, StateLayout, StateSegmentId, StateSegmentLifetime,
368 StateSegmentSpec, DEFAULT_STATE_SEGMENT_ID,
369};
370pub use weight_residency::{
371 DenseDiskStreamLoadOptions, DenseTransferSchedule, DenseTransferScheduleError,
372 ExecutionResidency, ExpertPass, LayerWeightResidency, LayerwiseLoadOptions,
373 LayerwiseModelMetadata, OrdinaryWeightResidency, ParameterBankAccess, ParameterBankKey,
374 ParameterBankLoadOptions, ParameterBankResidency, StaticUnitBindings, WeightResidency,
375 WeightResidencyPolicyError, DENSE_TRANSFER_WINDOW,
376};