Skip to main content

ferrum_interfaces/vnext/
standard_operations.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::num::NonZeroU32;
3
4use super::{
5    AliasPolicy, AttributeConstraint, AttributeId, AttributeSchema, AttributeSpec,
6    AttributeValueKind, CanonicalRational, CapabilityId, ContractVersion, DimensionConstraint,
7    ElementType, LayoutConstraint, OperationContract, OperationDescriptor, OperationId, OracleSpec,
8    ProfilePhase, ProviderRequirement, ResourcePresenceRequirement, ResourceRequirements,
9    TensorAccess, TensorContract, VNextError,
10};
11
12pub const TOKEN_EMBEDDING_OPERATION_ID: &str = "operation.token_embedding";
13pub const TOKEN_EMBEDDING_F16_CAPABILITY_ID: &str = "capability.operation.token_embedding.f16";
14pub const TOKEN_EMBEDDING_F32_MASTER_OPERATION_ID: &str = "operation.token_embedding.f32-master";
15pub const TOKEN_EMBEDDING_F32_MASTER_CAPABILITY_ID: &str =
16    "capability.operation.token_embedding.f32-master";
17pub const LAST_TOKEN_DENSE_LINEAR_OPERATION_ID: &str = "operation.last_token_dense_linear";
18pub const LAST_TOKEN_DENSE_LINEAR_F16_CAPABILITY_ID: &str =
19    "capability.operation.last_token_dense_linear.f16";
20pub const LAST_TOKEN_DENSE_LINEAR_F32_OPERATION_ID: &str = "operation.last_token_dense_linear.f32";
21pub const LAST_TOKEN_DENSE_LINEAR_F32_CAPABILITY_ID: &str =
22    "capability.operation.last_token_dense_linear.f32";
23pub const LAST_TOKEN_MASKED_ARGMAX_OPERATION_ID: &str = "operation.last_token_masked_argmax";
24pub const LAST_TOKEN_MASKED_ARGMAX_F16_CAPABILITY_ID: &str =
25    "capability.operation.last_token_masked_argmax.f16";
26pub const LAST_TOKEN_MASKED_ARGMAX_F32_OPERATION_ID: &str =
27    "operation.last_token_masked_argmax.f32";
28pub const LAST_TOKEN_MASKED_ARGMAX_F32_CAPABILITY_ID: &str =
29    "capability.operation.last_token_masked_argmax.f32";
30pub const RMS_NORM_OPERATION_ID: &str = "operation.rms_norm";
31pub const RMS_NORM_F16_CAPABILITY_ID: &str = "capability.operation.rms_norm.f16";
32pub const RMS_NORM_F32_TO_F16_OPERATION_ID: &str = "operation.rms_norm.f32-to-f16";
33pub const RMS_NORM_F32_TO_F16_CAPABILITY_ID: &str = "capability.operation.rms_norm.f32-to-f16";
34pub const RMS_NORM_F32_OPERATION_ID: &str = "operation.rms_norm.f32";
35pub const RMS_NORM_F32_CAPABILITY_ID: &str = "capability.operation.rms_norm.f32";
36pub const DENSE_LINEAR_OPERATION_ID: &str = "operation.dense_linear";
37pub const DENSE_LINEAR_F16_CAPABILITY_ID: &str = "capability.operation.dense_linear.f16";
38pub const DENSE_SWIGLU_OPERATION_ID: &str = "operation.dense_swiglu";
39pub const DENSE_SWIGLU_F16_CAPABILITY_ID: &str = "capability.operation.dense_swiglu.f16";
40pub const DENSE_GEGLU_TANH_OPERATION_ID: &str = "operation.dense_geglu_tanh";
41pub const DENSE_GEGLU_TANH_F16_CAPABILITY_ID: &str = "capability.operation.dense_geglu_tanh.f16";
42pub const CONSTANT_SCALE_OPERATION_ID: &str = "operation.constant_scale";
43pub const CONSTANT_SCALE_F16_CAPABILITY_ID: &str = "capability.operation.constant_scale.f16";
44pub const LOGIT_SOFTCAP_OPERATION_ID: &str = "operation.logit_softcap";
45pub const LOGIT_SOFTCAP_F16_CAPABILITY_ID: &str = "capability.operation.logit_softcap.f16";
46pub const ROUTED_SWIGLU_MOE_OPERATION_ID: &str = "operation.routed_swiglu_moe";
47pub const ROUTED_SWIGLU_MOE_F16_CAPABILITY_ID: &str = "capability.operation.routed_swiglu_moe.f16";
48pub const ROUTED_SHARED_SWIGLU_MOE_OPERATION_ID: &str = "operation.routed_shared_swiglu_moe";
49pub const ROUTED_SHARED_SWIGLU_MOE_F16_CAPABILITY_ID: &str =
50    "capability.operation.routed_shared_swiglu_moe.f16";
51pub const RESIDUAL_ADD_OPERATION_ID: &str = "operation.residual_add";
52pub const RESIDUAL_ADD_F16_CAPABILITY_ID: &str = "capability.operation.residual_add.f16";
53pub const RESIDUAL_ADD_F32_F16_OPERATION_ID: &str = "operation.residual_add.f32-f16";
54pub const RESIDUAL_ADD_F32_F16_CAPABILITY_ID: &str = "capability.operation.residual_add.f32-f16";
55pub const GATED_DELTA_RECURRENT_ATTENTION_OPERATION_ID: &str =
56    "operation.gated_delta_recurrent_attention";
57pub const GATED_DELTA_RECURRENT_ATTENTION_F16_CAPABILITY_ID: &str =
58    "capability.operation.gated_delta_recurrent_attention.f16";
59pub const GATED_DELTA_RECURRENT_ATTENTION_F32_MASTER_OPERATION_ID: &str =
60    "operation.gated_delta_recurrent_attention.f32-master";
61pub const GATED_DELTA_RECURRENT_ATTENTION_F32_MASTER_CAPABILITY_ID: &str =
62    "capability.operation.gated_delta_recurrent_attention.f32-master";
63pub const GATED_DELTA_EXECUTION_FORM_SELECTOR_VERSION: &str =
64    "gated-delta-execution-form-selector-v1";
65pub const CAUSAL_PAGED_ATTENTION_OPERATION_ID: &str = "operation.causal_paged_attention";
66pub const CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID: &str =
67    "capability.operation.causal_paged_attention.f16";
68pub const HYBRID_VNORM_CAUSAL_PAGED_ATTENTION_OPERATION_ID: &str =
69    "operation.hybrid_vnorm_causal_paged_attention";
70pub const HYBRID_VNORM_CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID: &str =
71    "capability.operation.hybrid_vnorm_causal_paged_attention.f16";
72pub const CAUSAL_PAGED_ATTENTION_F32_MASTER_OPERATION_ID: &str =
73    "operation.causal_paged_attention.f32-master";
74pub const CAUSAL_PAGED_ATTENTION_F32_MASTER_CAPABILITY_ID: &str =
75    "capability.operation.causal_paged_attention.f32-master";
76pub const CAUSAL_PAGED_ATTENTION_INT8_KV_OPERATION_ID: &str =
77    "operation.causal_paged_attention.int8-kv";
78pub const CAUSAL_PAGED_ATTENTION_INT8_KV_CAPABILITY_ID: &str =
79    "capability.operation.causal_paged_attention.f16.int8-kv";
80pub const CAUSAL_PAGED_ATTENTION_F32_MASTER_INT8_KV_OPERATION_ID: &str =
81    "operation.causal_paged_attention.f32-master.int8-kv";
82pub const CAUSAL_PAGED_ATTENTION_F32_MASTER_INT8_KV_CAPABILITY_ID: &str =
83    "capability.operation.causal_paged_attention.f32-master.int8-kv";
84pub const GPT_OSS_CAUSAL_PAGED_ATTENTION_OPERATION_ID: &str =
85    "operation.gpt_oss.causal_paged_attention";
86pub const GPT_OSS_CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID: &str =
87    "capability.operation.gpt_oss.causal_paged_attention.f16";
88pub const GPT_OSS_ROUTED_CLAMPED_SWIGLU_MOE_OPERATION_ID: &str =
89    "operation.gpt_oss.routed_clamped_swiglu_moe";
90pub const GPT_OSS_ROUTED_CLAMPED_SWIGLU_MOE_MXFP4_BF16_CAPABILITY_ID: &str =
91    "capability.operation.gpt_oss.routed_clamped_swiglu_moe.mxfp4_bf16";
92
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94pub enum GatedDeltaDecayParameterization {
95    LogRate,
96    NegativeRate,
97}
98
99impl GatedDeltaDecayParameterization {
100    pub const ALL: [Self; 2] = [Self::LogRate, Self::NegativeRate];
101
102    pub const fn as_str(self) -> &'static str {
103        match self {
104            Self::LogRate => "log_rate",
105            Self::NegativeRate => "negative_rate",
106        }
107    }
108
109    pub fn parse(value: &str) -> Option<Self> {
110        Self::ALL
111            .into_iter()
112            .find(|candidate| candidate.as_str() == value)
113    }
114}
115
116#[derive(Debug, Clone, Copy, PartialEq, Eq)]
117pub enum GatedDeltaValueHeadMapping {
118    GroupedByKeyHead,
119    InterleavedByKeyHead,
120}
121
122impl GatedDeltaValueHeadMapping {
123    pub const ALL: [Self; 2] = [Self::GroupedByKeyHead, Self::InterleavedByKeyHead];
124
125    pub const fn as_str(self) -> &'static str {
126        match self {
127            Self::GroupedByKeyHead => "grouped_by_key_head",
128            Self::InterleavedByKeyHead => "interleaved_by_key_head",
129        }
130    }
131
132    pub fn parse(value: &str) -> Option<Self> {
133        Self::ALL
134            .into_iter()
135            .find(|candidate| candidate.as_str() == value)
136    }
137}
138
139/// Physical gated-delta implementation available to one provider for an
140/// already-compatible operation shape. This capability is deliberately not a
141/// model attribute: the same immutable model plan can select a different form
142/// as the request work shape changes.
143#[derive(Debug, Clone, Copy, PartialEq, Eq)]
144pub struct GatedDeltaExecutionCapabilities {
145    chunked_scan: Option<GatedDeltaChunkedScanCapability>,
146}
147
148impl GatedDeltaExecutionCapabilities {
149    pub const fn recurrent_only() -> Self {
150        Self { chunked_scan: None }
151    }
152
153    pub fn with_chunked_scan(chunk_size: u32) -> Result<Self, VNextError> {
154        Ok(Self {
155            chunked_scan: Some(GatedDeltaChunkedScanCapability::new(chunk_size)?),
156        })
157    }
158
159    pub const fn chunked_scan(self) -> Option<GatedDeltaChunkedScanCapability> {
160        self.chunked_scan
161    }
162
163    /// Selects a physical form for one participant. Providers must first
164    /// remove capabilities that do not support the resolved dtype or shape.
165    pub fn select(
166        self,
167        token_count: u64,
168        preference: GatedDeltaExecutionPreference,
169    ) -> Result<GatedDeltaExecutionForm, VNextError> {
170        if token_count == 0 {
171            return Err(VNextError::InvalidExecutionPlan {
172                reason: "gated-delta execution requires at least one token".to_owned(),
173            });
174        }
175        match (preference, self.chunked_scan, token_count) {
176            (GatedDeltaExecutionPreference::ChunkedScan, Some(capability), 2..) => {
177                Ok(GatedDeltaExecutionForm::ChunkedScan(
178                    GatedDeltaChunkPlan::new(token_count, capability.chunk_size),
179                ))
180            }
181            _ => Ok(GatedDeltaExecutionForm::RecurrentScan),
182        }
183    }
184}
185
186#[derive(Debug, Clone, Copy, PartialEq, Eq)]
187pub struct GatedDeltaChunkedScanCapability {
188    chunk_size: NonZeroU32,
189}
190
191impl GatedDeltaChunkedScanCapability {
192    fn new(chunk_size: u32) -> Result<Self, VNextError> {
193        let chunk_size =
194            NonZeroU32::new(chunk_size).ok_or_else(|| VNextError::InvalidExecutionPlan {
195                reason: "gated-delta chunk size must be positive".to_owned(),
196            })?;
197        Ok(Self { chunk_size })
198    }
199
200    pub const fn chunk_size(self) -> u32 {
201        self.chunk_size.get()
202    }
203}
204
205/// Cost-model preference kept separate from physical support. A provider may
206/// derive it from calibrated crossover data and live batch topology without
207/// changing the immutable model plan.
208#[derive(Debug, Clone, Copy, PartialEq, Eq)]
209pub enum GatedDeltaExecutionPreference {
210    RecurrentScan,
211    ChunkedScan,
212}
213
214#[derive(Debug, Clone, Copy, PartialEq, Eq)]
215pub struct GatedDeltaChunkPlan {
216    token_count: u64,
217    chunk_size: NonZeroU32,
218    chunk_count: u64,
219    final_chunk_tokens: u32,
220}
221
222impl GatedDeltaChunkPlan {
223    fn new(token_count: u64, chunk_size: NonZeroU32) -> Self {
224        debug_assert!(token_count > 0);
225        let chunk_size_u64 = u64::from(chunk_size.get());
226        let chunk_count = ((token_count - 1) / chunk_size_u64) + 1;
227        let remainder = (token_count % chunk_size_u64) as u32;
228        Self {
229            token_count,
230            chunk_size,
231            chunk_count,
232            final_chunk_tokens: if remainder == 0 {
233                chunk_size.get()
234            } else {
235                remainder
236            },
237        }
238    }
239
240    pub const fn token_count(self) -> u64 {
241        self.token_count
242    }
243
244    pub const fn chunk_size(self) -> u32 {
245        self.chunk_size.get()
246    }
247
248    pub const fn chunk_count(self) -> u64 {
249        self.chunk_count
250    }
251
252    pub const fn final_chunk_tokens(self) -> u32 {
253        self.final_chunk_tokens
254    }
255}
256
257#[derive(Debug, Clone, Copy, PartialEq, Eq)]
258pub enum GatedDeltaExecutionForm {
259    RecurrentScan,
260    ChunkedScan(GatedDeltaChunkPlan),
261}
262
263impl GatedDeltaExecutionForm {
264    pub const fn as_str(self) -> &'static str {
265        match self {
266            Self::RecurrentScan => "recurrent_scan",
267            Self::ChunkedScan(_) => "chunked_scan",
268        }
269    }
270}
271
272/// One checked-in standard operation contract. Construction stays private so
273/// production registries cannot mutate a descriptor after a provider binds its
274/// fingerprint.
275pub struct StandardOperationContract {
276    descriptor: OperationDescriptor,
277}
278
279impl OperationContract for StandardOperationContract {
280    fn descriptor(&self) -> &OperationDescriptor {
281        &self.descriptor
282    }
283
284    fn validate_signature(
285        &self,
286        inputs: &[TensorContract],
287        outputs: &[TensorContract],
288    ) -> Result<(), VNextError> {
289        if inputs != self.descriptor.inputs || outputs != self.descriptor.outputs {
290            return Err(VNextError::InvalidExecutionPlan {
291                reason: format!(
292                    "operation `{}` signature differs from its standard contract",
293                    self.descriptor.id
294                ),
295            });
296        }
297        Ok(())
298    }
299}
300
301pub fn token_embedding_contract() -> Result<StandardOperationContract, VNextError> {
302    token_embedding_contract_with_output(
303        TOKEN_EMBEDDING_OPERATION_ID,
304        TOKEN_EMBEDDING_F16_CAPABILITY_ID,
305        ElementType::F16,
306    )
307}
308
309pub fn token_embedding_f32_master_contract() -> Result<StandardOperationContract, VNextError> {
310    token_embedding_contract_with_output(
311        TOKEN_EMBEDDING_F32_MASTER_OPERATION_ID,
312        TOKEN_EMBEDDING_F32_MASTER_CAPABILITY_ID,
313        ElementType::F32,
314    )
315}
316
317fn token_embedding_contract_with_output(
318    operation_id: &str,
319    capability_id: &str,
320    output_type: ElementType,
321) -> Result<StandardOperationContract, VNextError> {
322    let descriptor = OperationDescriptor {
323        id: OperationId::new(operation_id)?,
324        version: ContractVersion::new(1, 1),
325        inputs: vec![
326            contiguous_tensor(
327                vec![DimensionConstraint::Symbol("tokens".to_owned())],
328                [ElementType::U32],
329                TensorAccess::Read,
330            )?,
331            contiguous_tensor(
332                vec![
333                    DimensionConstraint::Symbol("vocab_size".to_owned()),
334                    DimensionConstraint::Symbol("hidden_size".to_owned()),
335                ],
336                [ElementType::F16],
337                TensorAccess::Read,
338            )?,
339        ],
340        outputs: vec![contiguous_tensor(
341            vec![
342                DimensionConstraint::Symbol("tokens".to_owned()),
343                DimensionConstraint::Symbol("hidden_size".to_owned()),
344            ],
345            [output_type],
346            TensorAccess::Write,
347        )?],
348        attributes: AttributeSchema::new(BTreeMap::from([
349            unsigned_attribute("hidden_size")?,
350            unsigned_attribute("vocab_size")?,
351        ]))?,
352        resources: ResourceRequirements {
353            minimum_value_alignment_bytes: 16,
354            // A transformed table needs bounded activation storage to restore
355            // each selected row without expanding the persistent weights.
356            scratch: ResourcePresenceRequirement::Optional,
357            binding: ResourcePresenceRequirement::Forbidden,
358            persistent: ResourcePresenceRequirement::Forbidden,
359        },
360        oracle: OracleSpec::Exact,
361        provider: ProviderRequirement {
362            minimum_version: ContractVersion::new(1, 0),
363            required_capabilities: BTreeSet::from([CapabilityId::new(capability_id)?]),
364        },
365        profile_phase: ProfilePhase::Forward,
366    };
367    descriptor.validate()?;
368    Ok(StandardOperationContract { descriptor })
369}
370
371/// Projects only the final row of a non-empty token-major tensor. Keeping this
372/// semantic fusion explicit prevents materializing prompt-length vocabulary
373/// logits while leaving providers free to use a pointer offset, row gather,
374/// or a fused kernel.
375pub fn last_token_dense_linear_contract() -> Result<StandardOperationContract, VNextError> {
376    last_token_dense_linear_contract_with_activation(
377        LAST_TOKEN_DENSE_LINEAR_OPERATION_ID,
378        ContractVersion::new(1, 1),
379        LAST_TOKEN_DENSE_LINEAR_F16_CAPABILITY_ID,
380        ElementType::F16,
381    )
382}
383
384pub fn last_token_dense_linear_f32_contract() -> Result<StandardOperationContract, VNextError> {
385    last_token_dense_linear_contract_with_activation(
386        LAST_TOKEN_DENSE_LINEAR_F32_OPERATION_ID,
387        ContractVersion::new(1, 0),
388        LAST_TOKEN_DENSE_LINEAR_F32_CAPABILITY_ID,
389        ElementType::F32,
390    )
391}
392
393fn last_token_dense_linear_contract_with_activation(
394    operation_id: &str,
395    version: ContractVersion,
396    capability_id: &str,
397    activation_type: ElementType,
398) -> Result<StandardOperationContract, VNextError> {
399    let descriptor = OperationDescriptor {
400        id: OperationId::new(operation_id)?,
401        version,
402        inputs: vec![
403            contiguous_tensor(
404                token_hidden_dimensions(),
405                [activation_type],
406                TensorAccess::Read,
407            )?,
408            contiguous_tensor(
409                vec![
410                    DimensionConstraint::Symbol("out_features".to_owned()),
411                    DimensionConstraint::Symbol("hidden_size".to_owned()),
412                ],
413                [ElementType::F16],
414                TensorAccess::Read,
415            )?,
416        ],
417        outputs: vec![contiguous_tensor(
418            vec![
419                DimensionConstraint::Exact(1),
420                DimensionConstraint::Symbol("out_features".to_owned()),
421            ],
422            [activation_type],
423            TensorAccess::Write,
424        )?],
425        attributes: AttributeSchema::new(BTreeMap::from([
426            unsigned_attribute("hidden_size")?,
427            unsigned_attribute("out_features")?,
428        ]))?,
429        resources: ResourceRequirements {
430            minimum_value_alignment_bytes: 16,
431            scratch: ResourcePresenceRequirement::Optional,
432            binding: ResourcePresenceRequirement::Forbidden,
433            persistent: ResourcePresenceRequirement::Forbidden,
434        },
435        oracle: f16_reference_tolerance()?,
436        provider: provider_requirement(capability_id, version)?,
437        profile_phase: ProfilePhase::Forward,
438    };
439    descriptor.validate()?;
440    Ok(StandardOperationContract { descriptor })
441}
442
443/// Selects one token from a final-position F16 logits row after applying an
444/// exact per-vocabulary validity mask and an optional sparse repetition
445/// penalty. Selection policy is carried by typed inputs so it remains visible
446/// to planning and cannot be hidden in backend flags. Semantic logits remain
447/// immutable; providers use invocation-scoped scratch for any penalized view.
448///
449/// The repetition token ids are unique and occupy
450/// `offsets[0]..offsets[1]` within the fixed-capacity input. A penalty of `1.0`
451/// or an empty range leaves logits unchanged.
452pub fn last_token_masked_argmax_contract() -> Result<StandardOperationContract, VNextError> {
453    last_token_masked_argmax_contract_with_logits(
454        LAST_TOKEN_MASKED_ARGMAX_OPERATION_ID,
455        ContractVersion::new(3, 0),
456        LAST_TOKEN_MASKED_ARGMAX_F16_CAPABILITY_ID,
457        ElementType::F16,
458    )
459}
460
461pub fn last_token_masked_argmax_f32_contract() -> Result<StandardOperationContract, VNextError> {
462    last_token_masked_argmax_contract_with_logits(
463        LAST_TOKEN_MASKED_ARGMAX_F32_OPERATION_ID,
464        ContractVersion::new(1, 0),
465        LAST_TOKEN_MASKED_ARGMAX_F32_CAPABILITY_ID,
466        ElementType::F32,
467    )
468}
469
470fn last_token_masked_argmax_contract_with_logits(
471    operation_id: &str,
472    version: ContractVersion,
473    capability_id: &str,
474    logits_type: ElementType,
475) -> Result<StandardOperationContract, VNextError> {
476    let descriptor = OperationDescriptor {
477        id: OperationId::new(operation_id)?,
478        version,
479        inputs: vec![
480            contiguous_tensor(
481                vec![
482                    DimensionConstraint::Exact(1),
483                    DimensionConstraint::Symbol("vocab_size".to_owned()),
484                ],
485                [logits_type],
486                TensorAccess::Read,
487            )?,
488            contiguous_tensor(
489                vec![DimensionConstraint::Symbol("vocab_size".to_owned())],
490                [ElementType::U8],
491                TensorAccess::Read,
492            )?,
493            contiguous_tensor(
494                vec![DimensionConstraint::Symbol(
495                    "repetition_capacity".to_owned(),
496                )],
497                [ElementType::U32],
498                TensorAccess::Read,
499            )?,
500            contiguous_tensor(
501                vec![DimensionConstraint::Exact(2)],
502                [ElementType::U32],
503                TensorAccess::Read,
504            )?,
505            contiguous_tensor(
506                vec![DimensionConstraint::Exact(1)],
507                [ElementType::F32],
508                TensorAccess::Read,
509            )?,
510        ],
511        outputs: vec![contiguous_tensor(
512            vec![DimensionConstraint::Exact(1)],
513            [ElementType::U32],
514            TensorAccess::Write,
515        )?],
516        attributes: AttributeSchema::new(BTreeMap::from([unsigned_attribute("vocab_size")?]))?,
517        resources: ResourceRequirements {
518            minimum_value_alignment_bytes: 16,
519            scratch: ResourcePresenceRequirement::Required,
520            binding: ResourcePresenceRequirement::Forbidden,
521            persistent: ResourcePresenceRequirement::Forbidden,
522        },
523        oracle: OracleSpec::Exact,
524        provider: provider_requirement(capability_id, version)?,
525        profile_phase: ProfilePhase::Forward,
526    };
527    descriptor.validate()?;
528    Ok(StandardOperationContract { descriptor })
529}
530
531pub fn rms_norm_contract() -> Result<StandardOperationContract, VNextError> {
532    rms_norm_contract_with_types(
533        RMS_NORM_OPERATION_ID,
534        RMS_NORM_F16_CAPABILITY_ID,
535        ElementType::F16,
536        ElementType::F16,
537    )
538}
539
540pub fn rms_norm_f32_to_f16_contract() -> Result<StandardOperationContract, VNextError> {
541    rms_norm_contract_with_types(
542        RMS_NORM_F32_TO_F16_OPERATION_ID,
543        RMS_NORM_F32_TO_F16_CAPABILITY_ID,
544        ElementType::F32,
545        ElementType::F16,
546    )
547}
548
549pub fn rms_norm_f32_contract() -> Result<StandardOperationContract, VNextError> {
550    rms_norm_contract_with_types(
551        RMS_NORM_F32_OPERATION_ID,
552        RMS_NORM_F32_CAPABILITY_ID,
553        ElementType::F32,
554        ElementType::F32,
555    )
556}
557
558fn rms_norm_contract_with_types(
559    operation_id: &str,
560    capability_id: &str,
561    input_type: ElementType,
562    output_type: ElementType,
563) -> Result<StandardOperationContract, VNextError> {
564    let descriptor = OperationDescriptor {
565        id: OperationId::new(operation_id)?,
566        version: ContractVersion::new(1, 0),
567        inputs: vec![
568            contiguous_tensor(token_hidden_dimensions(), [input_type], TensorAccess::Read)?,
569            contiguous_tensor(
570                vec![DimensionConstraint::Symbol("hidden_size".to_owned())],
571                [ElementType::F16],
572                TensorAccess::Read,
573            )?,
574        ],
575        outputs: vec![contiguous_tensor(
576            token_hidden_dimensions(),
577            [output_type],
578            TensorAccess::Write,
579        )?],
580        attributes: AttributeSchema::new(BTreeMap::from([
581            unsigned_attribute("hidden_size")?,
582            positive_epsilon_attribute("epsilon")?,
583        ]))?,
584        resources: no_auxiliary_resources(),
585        oracle: if output_type == ElementType::F32 {
586            f32_reference_tolerance()?
587        } else {
588            f16_reference_tolerance()?
589        },
590        provider: provider_requirement(capability_id, ContractVersion::new(1, 0))?,
591        profile_phase: ProfilePhase::Forward,
592    };
593    descriptor.validate()?;
594    Ok(StandardOperationContract { descriptor })
595}
596
597pub fn dense_linear_contract() -> Result<StandardOperationContract, VNextError> {
598    let descriptor = OperationDescriptor {
599        id: OperationId::new(DENSE_LINEAR_OPERATION_ID)?,
600        version: ContractVersion::new(1, 1),
601        inputs: vec![
602            contiguous_tensor(
603                vec![
604                    DimensionConstraint::Symbol("rows".to_owned()),
605                    DimensionConstraint::Symbol("in_features".to_owned()),
606                ],
607                [ElementType::F16],
608                TensorAccess::Read,
609            )?,
610            contiguous_tensor(
611                vec![
612                    DimensionConstraint::Symbol("out_features".to_owned()),
613                    DimensionConstraint::Symbol("in_features".to_owned()),
614                ],
615                [ElementType::F16],
616                TensorAccess::Read,
617            )?,
618        ],
619        outputs: vec![contiguous_tensor(
620            vec![
621                DimensionConstraint::Symbol("rows".to_owned()),
622                DimensionConstraint::Symbol("out_features".to_owned()),
623            ],
624            [ElementType::F16],
625            TensorAccess::Write,
626        )?],
627        attributes: AttributeSchema::new(BTreeMap::from([
628            unsigned_attribute("in_features")?,
629            unsigned_attribute("out_features")?,
630        ]))?,
631        resources: ResourceRequirements {
632            // Providers may transform the input coordinates before the dot
633            // product. Plain matrix providers still declare no workspace.
634            scratch: ResourcePresenceRequirement::Optional,
635            ..no_auxiliary_resources()
636        },
637        oracle: f16_reference_tolerance()?,
638        provider: provider_requirement(DENSE_LINEAR_F16_CAPABILITY_ID, ContractVersion::new(1, 0))?,
639        profile_phase: ProfilePhase::Forward,
640    };
641    descriptor.validate()?;
642    Ok(StandardOperationContract { descriptor })
643}
644
645pub fn dense_swiglu_contract() -> Result<StandardOperationContract, VNextError> {
646    let descriptor = OperationDescriptor {
647        id: OperationId::new(DENSE_SWIGLU_OPERATION_ID)?,
648        version: ContractVersion::new(1, 0),
649        inputs: vec![
650            contiguous_tensor(
651                token_hidden_dimensions(),
652                [ElementType::F16],
653                TensorAccess::Read,
654            )?,
655            contiguous_tensor(
656                packed_gate_up_dimensions(),
657                [ElementType::F16],
658                TensorAccess::Read,
659            )?,
660            contiguous_tensor(
661                hidden_intermediate_dimensions(),
662                [ElementType::F16],
663                TensorAccess::Read,
664            )?,
665        ],
666        outputs: vec![contiguous_tensor(
667            token_hidden_dimensions(),
668            [ElementType::F16],
669            TensorAccess::Write,
670        )?],
671        attributes: AttributeSchema::new(BTreeMap::from([
672            unsigned_attribute("hidden_size")?,
673            unsigned_attribute("intermediate_size")?,
674        ]))?,
675        resources: ResourceRequirements {
676            minimum_value_alignment_bytes: 16,
677            scratch: ResourcePresenceRequirement::Required,
678            binding: ResourcePresenceRequirement::Forbidden,
679            persistent: ResourcePresenceRequirement::Forbidden,
680        },
681        oracle: f16_reference_tolerance()?,
682        provider: provider_requirement(DENSE_SWIGLU_F16_CAPABILITY_ID, ContractVersion::new(1, 0))?,
683        profile_phase: ProfilePhase::Forward,
684    };
685    descriptor.validate()?;
686    Ok(StandardOperationContract { descriptor })
687}
688
689/// Dense GeGLU using the tanh approximation of GELU.
690///
691/// Gate and up projections remain independent logical weights at this
692/// boundary. Physical packing or quantization is a provider concern and must
693/// never be inferred from the operation signature.
694pub fn dense_geglu_tanh_contract() -> Result<StandardOperationContract, VNextError> {
695    let descriptor = OperationDescriptor {
696        id: OperationId::new(DENSE_GEGLU_TANH_OPERATION_ID)?,
697        version: ContractVersion::new(1, 0),
698        inputs: vec![
699            contiguous_tensor(
700                token_hidden_dimensions(),
701                [ElementType::F16],
702                TensorAccess::Read,
703            )?,
704            contiguous_tensor(
705                intermediate_hidden_dimensions(),
706                [ElementType::F16],
707                TensorAccess::Read,
708            )?,
709            contiguous_tensor(
710                intermediate_hidden_dimensions(),
711                [ElementType::F16],
712                TensorAccess::Read,
713            )?,
714            contiguous_tensor(
715                hidden_intermediate_dimensions(),
716                [ElementType::F16],
717                TensorAccess::Read,
718            )?,
719        ],
720        outputs: vec![contiguous_tensor(
721            token_hidden_dimensions(),
722            [ElementType::F16],
723            TensorAccess::Write,
724        )?],
725        attributes: AttributeSchema::new(BTreeMap::from([
726            unsigned_attribute("hidden_size")?,
727            unsigned_attribute("intermediate_size")?,
728        ]))?,
729        resources: ResourceRequirements {
730            minimum_value_alignment_bytes: 16,
731            scratch: ResourcePresenceRequirement::Required,
732            binding: ResourcePresenceRequirement::Forbidden,
733            persistent: ResourcePresenceRequirement::Forbidden,
734        },
735        oracle: f16_reference_tolerance()?,
736        provider: provider_requirement(
737            DENSE_GEGLU_TANH_F16_CAPABILITY_ID,
738            ContractVersion::new(1, 0),
739        )?,
740        profile_phase: ProfilePhase::Forward,
741    };
742    descriptor.validate()?;
743    Ok(StandardOperationContract { descriptor })
744}
745
746/// Multiplies a token-major F16 hidden tensor by one positive compile-time
747/// rational. The output must exactly alias the input so the operation cannot
748/// silently materialize an extra residual-sized buffer.
749pub fn constant_scale_contract() -> Result<StandardOperationContract, VNextError> {
750    let descriptor = OperationDescriptor {
751        id: OperationId::new(CONSTANT_SCALE_OPERATION_ID)?,
752        version: ContractVersion::new(1, 0),
753        inputs: vec![contiguous_tensor(
754            token_hidden_dimensions(),
755            [ElementType::F16],
756            TensorAccess::Read,
757        )?],
758        outputs: vec![contiguous_tensor_with_alias(
759            token_hidden_dimensions(),
760            [ElementType::F16],
761            TensorAccess::Write,
762            AliasPolicy::MustAlias { tensor_index: 0 },
763        )?],
764        attributes: AttributeSchema::new(BTreeMap::from([
765            unsigned_attribute("hidden_size")?,
766            positive_rational_attribute("scale")?,
767        ]))?,
768        resources: no_auxiliary_resources(),
769        oracle: f16_reference_tolerance()?,
770        provider: provider_requirement(
771            CONSTANT_SCALE_F16_CAPABILITY_ID,
772            ContractVersion::new(1, 0),
773        )?,
774        profile_phase: ProfilePhase::Forward,
775    };
776    descriptor.validate()?;
777    Ok(StandardOperationContract { descriptor })
778}
779
780/// Applies `cap * tanh(logit / cap)` to one final-position vocabulary row.
781/// The output is deliberately in-place so samplers consume the semantically
782/// capped logits without retaining a second vocabulary-sized allocation.
783pub fn logit_softcap_contract() -> Result<StandardOperationContract, VNextError> {
784    let dimensions = vec![
785        DimensionConstraint::Exact(1),
786        DimensionConstraint::Symbol("vocab_size".to_owned()),
787    ];
788    let descriptor = OperationDescriptor {
789        id: OperationId::new(LOGIT_SOFTCAP_OPERATION_ID)?,
790        version: ContractVersion::new(1, 0),
791        inputs: vec![contiguous_tensor(
792            dimensions.clone(),
793            [ElementType::F16],
794            TensorAccess::Read,
795        )?],
796        outputs: vec![contiguous_tensor_with_alias(
797            dimensions,
798            [ElementType::F16],
799            TensorAccess::Write,
800            AliasPolicy::MustAlias { tensor_index: 0 },
801        )?],
802        attributes: AttributeSchema::new(BTreeMap::from([
803            unsigned_attribute("vocab_size")?,
804            positive_rational_attribute("cap")?,
805        ]))?,
806        resources: no_auxiliary_resources(),
807        oracle: f16_reference_tolerance()?,
808        provider: provider_requirement(
809            LOGIT_SOFTCAP_F16_CAPABILITY_ID,
810            ContractVersion::new(1, 0),
811        )?,
812        profile_phase: ProfilePhase::Forward,
813    };
814    descriptor.validate()?;
815    Ok(StandardOperationContract { descriptor })
816}
817
818/// A routed SwiGLU expert set plus one sigmoid-gated shared SwiGLU expert.
819///
820/// The operation boundary intentionally owns routing, routed expert execution,
821/// shared expert execution, and the final sum. Providers can choose a
822/// monolithic kernel, overlap the shared path with routed experts, or use a
823/// decomposed fallback without changing the immutable model program. Weight
824/// ordinals and logical stack shapes are part of the stable contract; physical
825/// quantization and expert placement remain weight/provider concerns.
826pub fn routed_shared_swiglu_moe_contract() -> Result<StandardOperationContract, VNextError> {
827    let descriptor = OperationDescriptor {
828        id: OperationId::new(ROUTED_SHARED_SWIGLU_MOE_OPERATION_ID)?,
829        version: ContractVersion::new(1, 0),
830        inputs: vec![
831            contiguous_tensor(
832                token_hidden_dimensions(),
833                [ElementType::F16],
834                TensorAccess::Read,
835            )?,
836            contiguous_tensor(
837                vec![symbol("expert_count"), symbol("hidden_size")],
838                [ElementType::F16],
839                TensorAccess::Read,
840            )?,
841            contiguous_tensor(
842                routed_expert_gate_up_dimensions(),
843                [ElementType::F16],
844                TensorAccess::Read,
845            )?,
846            contiguous_tensor(
847                routed_expert_down_dimensions(),
848                [ElementType::F16],
849                TensorAccess::Read,
850            )?,
851            contiguous_tensor(
852                vec![exact(1), symbol("hidden_size")],
853                [ElementType::F16],
854                TensorAccess::Read,
855            )?,
856            contiguous_tensor(
857                shared_expert_gate_up_dimensions(),
858                [ElementType::F16],
859                TensorAccess::Read,
860            )?,
861            contiguous_tensor(
862                shared_expert_down_dimensions(),
863                [ElementType::F16],
864                TensorAccess::Read,
865            )?,
866        ],
867        outputs: vec![contiguous_tensor(
868            token_hidden_dimensions(),
869            [ElementType::F16],
870            TensorAccess::Write,
871        )?],
872        attributes: AttributeSchema::new(BTreeMap::from([
873            unsigned_attribute("hidden_size")?,
874            unsigned_attribute("expert_count")?,
875            unsigned_attribute("experts_per_token")?,
876            unsigned_attribute("routed_intermediate_size")?,
877            unsigned_attribute("shared_intermediate_size")?,
878            unconstrained_bool_attribute("normalize_topk")?,
879        ]))?,
880        resources: ResourceRequirements {
881            minimum_value_alignment_bytes: 16,
882            scratch: ResourcePresenceRequirement::Required,
883            binding: ResourcePresenceRequirement::Forbidden,
884            persistent: ResourcePresenceRequirement::Forbidden,
885        },
886        oracle: f16_reference_tolerance()?,
887        provider: provider_requirement(
888            ROUTED_SHARED_SWIGLU_MOE_F16_CAPABILITY_ID,
889            ContractVersion::new(1, 0),
890        )?,
891        profile_phase: ProfilePhase::Forward,
892    };
893    descriptor.validate()?;
894    Ok(StandardOperationContract { descriptor })
895}
896
897/// A top-K routed SwiGLU expert set without a shared expert branch.
898///
899/// Routing, expert execution, weighted reduction, and their scratch lifetime
900/// form one stable operation boundary. This is intentionally separate from
901/// [`routed_shared_swiglu_moe_contract`]: a provider must never synthesize
902/// shared-expert weights or execute extra shared work for routed-only model
903/// families.
904pub fn routed_swiglu_moe_contract() -> Result<StandardOperationContract, VNextError> {
905    let descriptor = OperationDescriptor {
906        id: OperationId::new(ROUTED_SWIGLU_MOE_OPERATION_ID)?,
907        version: ContractVersion::new(1, 0),
908        inputs: vec![
909            contiguous_tensor(
910                token_hidden_dimensions(),
911                [ElementType::F16],
912                TensorAccess::Read,
913            )?,
914            contiguous_tensor(
915                vec![symbol("expert_count"), symbol("hidden_size")],
916                [ElementType::F16],
917                TensorAccess::Read,
918            )?,
919            contiguous_tensor(
920                routed_expert_gate_up_dimensions(),
921                [ElementType::F16],
922                TensorAccess::Read,
923            )?,
924            contiguous_tensor(
925                routed_expert_down_dimensions(),
926                [ElementType::F16],
927                TensorAccess::Read,
928            )?,
929        ],
930        outputs: vec![contiguous_tensor(
931            token_hidden_dimensions(),
932            [ElementType::F16],
933            TensorAccess::Write,
934        )?],
935        attributes: AttributeSchema::new(BTreeMap::from([
936            unsigned_attribute("hidden_size")?,
937            unsigned_attribute("expert_count")?,
938            unsigned_attribute("experts_per_token")?,
939            unsigned_attribute("routed_intermediate_size")?,
940            unconstrained_bool_attribute("normalize_topk")?,
941        ]))?,
942        resources: ResourceRequirements {
943            minimum_value_alignment_bytes: 16,
944            scratch: ResourcePresenceRequirement::Required,
945            binding: ResourcePresenceRequirement::Forbidden,
946            persistent: ResourcePresenceRequirement::Forbidden,
947        },
948        oracle: f16_reference_tolerance()?,
949        provider: provider_requirement(
950            ROUTED_SWIGLU_MOE_F16_CAPABILITY_ID,
951            ContractVersion::new(1, 0),
952        )?,
953        profile_phase: ProfilePhase::Forward,
954    };
955    descriptor.validate()?;
956    Ok(StandardOperationContract { descriptor })
957}
958
959/// GPT-OSS top-K routed experts with the model's clamped, interleaved SwiGLU.
960///
961/// Expert matrices are logical BF16 tensors at this boundary. Their checkpoint
962/// packing, sidecar scales, and any execution repack remain exclusively in the
963/// weight materializer/provider contracts. `gate_up_features` is explicit
964/// because the generic tensor contract has no derived-dimension expression;
965/// callers and providers must bind it to exactly twice `intermediate_size`.
966/// With `gate_up_interleaved=true`, even rows are gate rows and odd rows are up
967/// rows in the logical `[E, 2*I, H]` tensor; no implicit reshape is permitted.
968pub fn gpt_oss_routed_clamped_swiglu_moe_contract() -> Result<StandardOperationContract, VNextError>
969{
970    let descriptor = OperationDescriptor {
971        id: OperationId::new(GPT_OSS_ROUTED_CLAMPED_SWIGLU_MOE_OPERATION_ID)?,
972        version: ContractVersion::new(1, 0),
973        inputs: vec![
974            contiguous_tensor(
975                token_hidden_dimensions(),
976                [ElementType::F16],
977                TensorAccess::Read,
978            )?,
979            contiguous_tensor(
980                vec![symbol("expert_count"), symbol("hidden_size")],
981                [ElementType::Bf16],
982                TensorAccess::Read,
983            )?,
984            contiguous_tensor(
985                vec![symbol("expert_count")],
986                [ElementType::Bf16],
987                TensorAccess::Read,
988            )?,
989            contiguous_tensor(
990                vec![
991                    symbol("expert_count"),
992                    symbol("gate_up_features"),
993                    symbol("hidden_size"),
994                ],
995                [ElementType::Bf16],
996                TensorAccess::Read,
997            )?,
998            contiguous_tensor(
999                vec![symbol("expert_count"), symbol("gate_up_features")],
1000                [ElementType::Bf16],
1001                TensorAccess::Read,
1002            )?,
1003            contiguous_tensor(
1004                vec![
1005                    symbol("expert_count"),
1006                    symbol("hidden_size"),
1007                    symbol("intermediate_size"),
1008                ],
1009                [ElementType::Bf16],
1010                TensorAccess::Read,
1011            )?,
1012            contiguous_tensor(
1013                vec![symbol("expert_count"), symbol("hidden_size")],
1014                [ElementType::Bf16],
1015                TensorAccess::Read,
1016            )?,
1017        ],
1018        outputs: vec![contiguous_tensor(
1019            token_hidden_dimensions(),
1020            [ElementType::F16],
1021            TensorAccess::Write,
1022        )?],
1023        attributes: AttributeSchema::new(BTreeMap::from([
1024            unsigned_attribute("hidden_size")?,
1025            unsigned_attribute("expert_count")?,
1026            unsigned_attribute("experts_per_token")?,
1027            unsigned_attribute("intermediate_size")?,
1028            unsigned_attribute("gate_up_features")?,
1029            true_bool_attribute("normalize_topk")?,
1030            exact_rational_attribute("swiglu_limit", 7, 1)?,
1031            true_bool_attribute("gate_up_interleaved")?,
1032            true_bool_attribute("down_bias_before_route_reduction")?,
1033        ]))?,
1034        resources: ResourceRequirements {
1035            minimum_value_alignment_bytes: 16,
1036            scratch: ResourcePresenceRequirement::Required,
1037            binding: ResourcePresenceRequirement::Forbidden,
1038            persistent: ResourcePresenceRequirement::Forbidden,
1039        },
1040        oracle: f16_reference_tolerance()?,
1041        provider: provider_requirement(
1042            GPT_OSS_ROUTED_CLAMPED_SWIGLU_MOE_MXFP4_BF16_CAPABILITY_ID,
1043            ContractVersion::new(1, 0),
1044        )?,
1045        profile_phase: ProfilePhase::Forward,
1046    };
1047    descriptor.validate()?;
1048    Ok(StandardOperationContract { descriptor })
1049}
1050
1051pub fn residual_add_contract() -> Result<StandardOperationContract, VNextError> {
1052    residual_add_contract_with_types(
1053        RESIDUAL_ADD_OPERATION_ID,
1054        RESIDUAL_ADD_F16_CAPABILITY_ID,
1055        ElementType::F16,
1056        ElementType::F16,
1057        ElementType::F16,
1058    )
1059}
1060
1061pub fn residual_add_f32_f16_contract() -> Result<StandardOperationContract, VNextError> {
1062    residual_add_contract_with_types(
1063        RESIDUAL_ADD_F32_F16_OPERATION_ID,
1064        RESIDUAL_ADD_F32_F16_CAPABILITY_ID,
1065        ElementType::F32,
1066        ElementType::F16,
1067        ElementType::F32,
1068    )
1069}
1070
1071fn residual_add_contract_with_types(
1072    operation_id: &str,
1073    capability_id: &str,
1074    left_type: ElementType,
1075    right_type: ElementType,
1076    output_type: ElementType,
1077) -> Result<StandardOperationContract, VNextError> {
1078    let descriptor = OperationDescriptor {
1079        id: OperationId::new(operation_id)?,
1080        version: ContractVersion::new(1, 0),
1081        inputs: vec![
1082            contiguous_tensor(token_hidden_dimensions(), [left_type], TensorAccess::Read)?,
1083            contiguous_tensor(token_hidden_dimensions(), [right_type], TensorAccess::Read)?,
1084        ],
1085        outputs: vec![contiguous_tensor_with_alias(
1086            token_hidden_dimensions(),
1087            [output_type],
1088            TensorAccess::Write,
1089            AliasPolicy::MayAlias { tensor_index: 0 },
1090        )?],
1091        attributes: AttributeSchema::new(BTreeMap::from([unsigned_attribute("hidden_size")?]))?,
1092        resources: no_auxiliary_resources(),
1093        oracle: if output_type == ElementType::F32 {
1094            OracleSpec::Exact
1095        } else {
1096            f16_reference_tolerance()?
1097        },
1098        provider: provider_requirement(capability_id, ContractVersion::new(1, 0))?,
1099        profile_phase: ProfilePhase::Forward,
1100    };
1101    descriptor.validate()?;
1102    Ok(StandardOperationContract { descriptor })
1103}
1104
1105/// Gated DeltaNet mixer including input normalization, projections, recurrent
1106/// convolution/Delta state update, gated normalization, output projection, and
1107/// the attention residual. Weight ordinals are part of the stable contract.
1108pub fn gated_delta_recurrent_attention_contract() -> Result<StandardOperationContract, VNextError> {
1109    gated_delta_recurrent_attention_contract_with_hidden(
1110        GATED_DELTA_RECURRENT_ATTENTION_OPERATION_ID,
1111        ContractVersion::new(6, 0),
1112        GATED_DELTA_RECURRENT_ATTENTION_F16_CAPABILITY_ID,
1113        ElementType::F16,
1114    )
1115}
1116
1117pub fn gated_delta_recurrent_attention_f32_master_contract(
1118) -> Result<StandardOperationContract, VNextError> {
1119    gated_delta_recurrent_attention_contract_with_hidden(
1120        GATED_DELTA_RECURRENT_ATTENTION_F32_MASTER_OPERATION_ID,
1121        ContractVersion::new(1, 0),
1122        GATED_DELTA_RECURRENT_ATTENTION_F32_MASTER_CAPABILITY_ID,
1123        ElementType::F32,
1124    )
1125}
1126
1127fn gated_delta_recurrent_attention_contract_with_hidden(
1128    operation_id: &str,
1129    version: ContractVersion,
1130    capability_id: &str,
1131    hidden_type: ElementType,
1132) -> Result<StandardOperationContract, VNextError> {
1133    let descriptor = OperationDescriptor {
1134        id: OperationId::new(operation_id)?,
1135        version,
1136        inputs: vec![
1137            contiguous_tensor(token_hidden_dimensions(), [hidden_type], TensorAccess::Read)?,
1138            contiguous_tensor(
1139                vec![symbol("hidden_size")],
1140                [ElementType::F16],
1141                TensorAccess::Read,
1142            )?,
1143            contiguous_tensor(
1144                vec![symbol("qkvzba_features"), symbol("hidden_size")],
1145                [ElementType::F16],
1146                TensorAccess::Read,
1147            )?,
1148            contiguous_tensor(
1149                vec![symbol("qkv_features"), symbol("conv_kernel")],
1150                [ElementType::F16],
1151                TensorAccess::Read,
1152            )?,
1153            contiguous_tensor(
1154                vec![symbol("value_heads")],
1155                [ElementType::F32],
1156                TensorAccess::Read,
1157            )?,
1158            contiguous_tensor(
1159                vec![symbol("value_heads")],
1160                [ElementType::F32],
1161                TensorAccess::Read,
1162            )?,
1163            contiguous_tensor(
1164                vec![symbol("value_head_dim")],
1165                [ElementType::F32],
1166                TensorAccess::Read,
1167            )?,
1168            contiguous_tensor(
1169                vec![symbol("hidden_size"), symbol("value_features")],
1170                [ElementType::F16],
1171                TensorAccess::Read,
1172            )?,
1173            contiguous_tensor(
1174                vec![symbol("qkv_features"), symbol("conv_state_width")],
1175                [ElementType::F16],
1176                TensorAccess::ReadWrite,
1177            )?,
1178            contiguous_tensor(
1179                vec![
1180                    symbol("value_heads"),
1181                    symbol("value_head_dim"),
1182                    symbol("key_head_dim"),
1183                ],
1184                [ElementType::F32],
1185                TensorAccess::ReadWrite,
1186            )?,
1187        ],
1188        outputs: vec![contiguous_tensor_with_alias(
1189            token_hidden_dimensions(),
1190            [hidden_type],
1191            TensorAccess::Write,
1192            AliasPolicy::MayAlias { tensor_index: 0 },
1193        )?],
1194        attributes: AttributeSchema::new(BTreeMap::from([
1195            unsigned_attribute("hidden_size")?,
1196            unsigned_attribute("key_heads")?,
1197            unsigned_attribute("value_heads")?,
1198            unsigned_attribute("key_head_dim")?,
1199            unsigned_attribute("value_head_dim")?,
1200            unsigned_attribute("qkv_features")?,
1201            unsigned_attribute("value_features")?,
1202            unsigned_attribute("qkvz_features")?,
1203            unsigned_attribute("ba_features")?,
1204            unsigned_attribute("qkvzba_features")?,
1205            unsigned_attribute("conv_kernel")?,
1206            unsigned_attribute("conv_state_width")?,
1207            positive_epsilon_attribute("epsilon")?,
1208            nonnegative_unsigned_attribute("layer_index")?,
1209            text_choices_attribute(
1210                "decay_parameterization",
1211                GatedDeltaDecayParameterization::ALL.map(|value| value.as_str()),
1212            )?,
1213            text_choices_attribute(
1214                "value_head_mapping",
1215                GatedDeltaValueHeadMapping::ALL.map(|value| value.as_str()),
1216            )?,
1217        ]))?,
1218        resources: attention_resources(),
1219        oracle: f16_reference_tolerance()?,
1220        provider: provider_requirement(capability_id, version)?,
1221        profile_phase: ProfilePhase::Forward,
1222    };
1223    descriptor.validate()?;
1224    Ok(StandardOperationContract { descriptor })
1225}
1226
1227/// Dense causal attention including input normalization, Q/K normalization,
1228/// RoPE, KV update, attention, optional output gate, output projection, and
1229/// the attention residual. KV physical paging remains a provider concern.
1230pub fn causal_paged_attention_contract() -> Result<StandardOperationContract, VNextError> {
1231    causal_paged_attention_contract_with_hidden(
1232        CAUSAL_PAGED_ATTENTION_OPERATION_ID,
1233        ContractVersion::new(2, 0),
1234        CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID,
1235        ElementType::F16,
1236        false,
1237    )
1238}
1239
1240pub fn causal_paged_attention_f32_master_contract() -> Result<StandardOperationContract, VNextError>
1241{
1242    causal_paged_attention_contract_with_hidden(
1243        CAUSAL_PAGED_ATTENTION_F32_MASTER_OPERATION_ID,
1244        ContractVersion::new(1, 0),
1245        CAUSAL_PAGED_ATTENTION_F32_MASTER_CAPABILITY_ID,
1246        ElementType::F32,
1247        false,
1248    )
1249}
1250
1251/// Causal attention with the existing F16 activation arithmetic and versioned
1252/// INT8 KV storage. Input 8 is I8 [2, KV heads, head dimension]; input 9 is F32
1253/// [2, KV heads]. Each absolute token position owns one tensor of each kind.
1254///
1255/// K and V independently round to the existing F16 KV write boundary, then
1256/// convert to F32. For each token/head vector, scale = max(abs(x)) / 127;
1257/// quantized = clamp(round_ties_even(x / stored_scale), -127, 127). Zero vectors
1258/// use scale 1 and payload 0. Non-finite input fails the step; it must not become
1259/// a reusable completed boundary. Every attention read, including the current
1260/// prefill chunk, consumes q * scale with the original arithmetic boundaries.
1261/// Providers may use bounded tiles, never an unaccounted full-history F16 copy.
1262pub fn causal_paged_attention_int8_kv_contract() -> Result<StandardOperationContract, VNextError> {
1263    causal_paged_attention_contract_with_hidden(
1264        CAUSAL_PAGED_ATTENTION_INT8_KV_OPERATION_ID,
1265        ContractVersion::new(1, 0),
1266        CAUSAL_PAGED_ATTENTION_INT8_KV_CAPABILITY_ID,
1267        ElementType::F16,
1268        true,
1269    )
1270}
1271
1272/// The same versioned INT8 KV ABI with the existing F32-master hidden boundary.
1273pub fn causal_paged_attention_f32_master_int8_kv_contract(
1274) -> Result<StandardOperationContract, VNextError> {
1275    causal_paged_attention_contract_with_hidden(
1276        CAUSAL_PAGED_ATTENTION_F32_MASTER_INT8_KV_OPERATION_ID,
1277        ContractVersion::new(1, 0),
1278        CAUSAL_PAGED_ATTENTION_F32_MASTER_INT8_KV_CAPABILITY_ID,
1279        ElementType::F32,
1280        true,
1281    )
1282}
1283
1284fn causal_paged_attention_contract_with_hidden(
1285    operation_id: &str,
1286    version: ContractVersion,
1287    capability_id: &str,
1288    hidden_type: ElementType,
1289    int8_kv: bool,
1290) -> Result<StandardOperationContract, VNextError> {
1291    let mut descriptor = OperationDescriptor {
1292        id: OperationId::new(operation_id)?,
1293        version,
1294        inputs: vec![
1295            contiguous_tensor(token_hidden_dimensions(), [hidden_type], TensorAccess::Read)?,
1296            contiguous_tensor(
1297                vec![symbol("hidden_size")],
1298                [ElementType::F16],
1299                TensorAccess::Read,
1300            )?,
1301            contiguous_tensor(
1302                vec![symbol("query_projection_features"), symbol("hidden_size")],
1303                [ElementType::F16],
1304                TensorAccess::Read,
1305            )?,
1306            contiguous_tensor(
1307                vec![symbol("kv_features"), symbol("hidden_size")],
1308                [ElementType::F16],
1309                TensorAccess::Read,
1310            )?,
1311            contiguous_tensor(
1312                vec![symbol("kv_features"), symbol("hidden_size")],
1313                [ElementType::F16],
1314                TensorAccess::Read,
1315            )?,
1316            contiguous_tensor(
1317                vec![symbol("hidden_size"), symbol("query_features")],
1318                [ElementType::F16],
1319                TensorAccess::Read,
1320            )?,
1321            contiguous_tensor(
1322                vec![symbol("head_dim")],
1323                [ElementType::F16],
1324                TensorAccess::Read,
1325            )?,
1326            contiguous_tensor(
1327                vec![symbol("head_dim")],
1328                [ElementType::F16],
1329                TensorAccess::Read,
1330            )?,
1331            contiguous_tensor(
1332                vec![exact(2), symbol("key_value_heads"), symbol("head_dim")],
1333                [if int8_kv {
1334                    ElementType::I8
1335                } else {
1336                    ElementType::F16
1337                }],
1338                TensorAccess::ReadWrite,
1339            )?,
1340        ],
1341        outputs: vec![contiguous_tensor_with_alias(
1342            token_hidden_dimensions(),
1343            [hidden_type],
1344            TensorAccess::Write,
1345            AliasPolicy::MayAlias { tensor_index: 0 },
1346        )?],
1347        attributes: AttributeSchema::new(BTreeMap::from([
1348            unsigned_attribute("hidden_size")?,
1349            unsigned_attribute("query_heads")?,
1350            unsigned_attribute("key_value_heads")?,
1351            unsigned_attribute("head_dim")?,
1352            unsigned_attribute("query_features")?,
1353            unsigned_attribute("query_projection_features")?,
1354            unsigned_attribute("kv_features")?,
1355            unsigned_attribute("rope_dim")?,
1356            unsigned_attribute("maximum_context_tokens")?,
1357            positive_rational_attribute("rope_theta")?,
1358            unconstrained_bool_attribute("rope_interleaved")?,
1359            unconstrained_bool_attribute("output_gate")?,
1360            true_bool_attribute("causal")?,
1361            positive_epsilon_attribute("epsilon")?,
1362            nonnegative_unsigned_attribute("layer_index")?,
1363        ]))?,
1364        resources: causal_attention_resources(),
1365        oracle: f16_reference_tolerance()?,
1366        provider: provider_requirement(capability_id, version)?,
1367        profile_phase: ProfilePhase::Forward,
1368    };
1369    if int8_kv {
1370        descriptor.inputs.push(contiguous_tensor(
1371            vec![exact(2), symbol("key_value_heads")],
1372            [ElementType::F32],
1373            TensorAccess::ReadWrite,
1374        )?);
1375    }
1376    descriptor.validate()?;
1377    Ok(StandardOperationContract { descriptor })
1378}
1379
1380/// Hybrid causal attention with value normalization and optional K-as-V.
1381///
1382/// In addition to the shared causal attention pipeline, this contract makes
1383/// the hybrid-layer semantics explicit: the active rotary width and its
1384/// frequency denominator are independent, attention uses a typed scale,
1385/// local layers carry a sliding window, values use weightless RMSNorm, full
1386/// layers may bind K as V, and post-attention RMSNorm is applied before the
1387/// residual is added.
1388pub fn hybrid_vnorm_causal_paged_attention_contract(
1389) -> Result<StandardOperationContract, VNextError> {
1390    let descriptor = OperationDescriptor {
1391        id: OperationId::new(HYBRID_VNORM_CAUSAL_PAGED_ATTENTION_OPERATION_ID)?,
1392        version: ContractVersion::new(1, 0),
1393        inputs: vec![
1394            contiguous_tensor(
1395                token_hidden_dimensions(),
1396                [ElementType::F16],
1397                TensorAccess::Read,
1398            )?,
1399            contiguous_tensor(
1400                vec![symbol("hidden_size")],
1401                [ElementType::F16],
1402                TensorAccess::Read,
1403            )?,
1404            contiguous_tensor(
1405                vec![symbol("query_projection_features"), symbol("hidden_size")],
1406                [ElementType::F16],
1407                TensorAccess::Read,
1408            )?,
1409            contiguous_tensor(
1410                vec![symbol("kv_features"), symbol("hidden_size")],
1411                [ElementType::F16],
1412                TensorAccess::Read,
1413            )?,
1414            contiguous_tensor(
1415                vec![symbol("kv_features"), symbol("hidden_size")],
1416                [ElementType::F16],
1417                TensorAccess::Read,
1418            )?,
1419            contiguous_tensor(
1420                vec![symbol("hidden_size"), symbol("query_features")],
1421                [ElementType::F16],
1422                TensorAccess::Read,
1423            )?,
1424            contiguous_tensor(
1425                vec![symbol("head_dim")],
1426                [ElementType::F16],
1427                TensorAccess::Read,
1428            )?,
1429            contiguous_tensor(
1430                vec![symbol("head_dim")],
1431                [ElementType::F16],
1432                TensorAccess::Read,
1433            )?,
1434            contiguous_tensor(
1435                vec![exact(2), symbol("key_value_heads"), symbol("head_dim")],
1436                [ElementType::F16],
1437                TensorAccess::ReadWrite,
1438            )?,
1439            contiguous_tensor(
1440                vec![symbol("hidden_size")],
1441                [ElementType::F16],
1442                TensorAccess::Read,
1443            )?,
1444        ],
1445        outputs: vec![contiguous_tensor_with_alias(
1446            token_hidden_dimensions(),
1447            [ElementType::F16],
1448            TensorAccess::Write,
1449            AliasPolicy::MayAlias { tensor_index: 0 },
1450        )?],
1451        attributes: AttributeSchema::new(BTreeMap::from([
1452            unsigned_attribute("hidden_size")?,
1453            unsigned_attribute("query_heads")?,
1454            unsigned_attribute("key_value_heads")?,
1455            unsigned_attribute("head_dim")?,
1456            unsigned_attribute("query_features")?,
1457            unsigned_attribute("query_projection_features")?,
1458            unsigned_attribute("kv_features")?,
1459            unsigned_attribute("rope_dim")?,
1460            unsigned_attribute("rope_frequency_denominator")?,
1461            unsigned_attribute("maximum_context_tokens")?,
1462            positive_rational_attribute("rope_theta")?,
1463            unconstrained_bool_attribute("rope_interleaved")?,
1464            positive_rational_attribute("attention_scale")?,
1465            nonnegative_unsigned_attribute("sliding_window_tokens")?,
1466            true_bool_attribute("value_rms_norm")?,
1467            unconstrained_bool_attribute("attention_k_eq_v")?,
1468            true_bool_attribute("causal")?,
1469            positive_epsilon_attribute("epsilon")?,
1470            nonnegative_unsigned_attribute("layer_index")?,
1471        ]))?,
1472        resources: causal_attention_resources(),
1473        oracle: f16_reference_tolerance()?,
1474        provider: provider_requirement(
1475            HYBRID_VNORM_CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID,
1476            ContractVersion::new(1, 0),
1477        )?,
1478        profile_phase: ProfilePhase::Forward,
1479    };
1480    descriptor.validate()?;
1481    Ok(StandardOperationContract { descriptor })
1482}
1483
1484/// GPT-OSS causal attention including input normalization, biased Q/K/V/O
1485/// projections, per-query-head attention sinks, YaRN RoPE, KV update, output
1486/// projection, and the attention residual. A zero `sliding_window_tokens`
1487/// selects full causal attention; a positive value selects the typed local
1488/// window. KV paging and kernel fusion remain provider concerns.
1489pub fn gpt_oss_causal_paged_attention_contract() -> Result<StandardOperationContract, VNextError> {
1490    let descriptor = OperationDescriptor {
1491        id: OperationId::new(GPT_OSS_CAUSAL_PAGED_ATTENTION_OPERATION_ID)?,
1492        version: ContractVersion::new(1, 0),
1493        inputs: vec![
1494            contiguous_tensor(
1495                token_hidden_dimensions(),
1496                [ElementType::F16],
1497                TensorAccess::Read,
1498            )?,
1499            contiguous_tensor(
1500                vec![symbol("hidden_size")],
1501                [ElementType::F16],
1502                TensorAccess::Read,
1503            )?,
1504            contiguous_tensor(
1505                vec![symbol("query_features"), symbol("hidden_size")],
1506                [ElementType::F16],
1507                TensorAccess::Read,
1508            )?,
1509            contiguous_tensor(
1510                vec![symbol("kv_features"), symbol("hidden_size")],
1511                [ElementType::F16],
1512                TensorAccess::Read,
1513            )?,
1514            contiguous_tensor(
1515                vec![symbol("kv_features"), symbol("hidden_size")],
1516                [ElementType::F16],
1517                TensorAccess::Read,
1518            )?,
1519            contiguous_tensor(
1520                vec![symbol("hidden_size"), symbol("query_features")],
1521                [ElementType::F16],
1522                TensorAccess::Read,
1523            )?,
1524            contiguous_tensor(
1525                vec![symbol("query_features")],
1526                [ElementType::F16],
1527                TensorAccess::Read,
1528            )?,
1529            contiguous_tensor(
1530                vec![symbol("kv_features")],
1531                [ElementType::F16],
1532                TensorAccess::Read,
1533            )?,
1534            contiguous_tensor(
1535                vec![symbol("kv_features")],
1536                [ElementType::F16],
1537                TensorAccess::Read,
1538            )?,
1539            contiguous_tensor(
1540                vec![symbol("hidden_size")],
1541                [ElementType::F16],
1542                TensorAccess::Read,
1543            )?,
1544            contiguous_tensor(
1545                vec![symbol("query_heads")],
1546                [ElementType::F16],
1547                TensorAccess::Read,
1548            )?,
1549            contiguous_tensor(
1550                vec![exact(2), symbol("kv_heads"), symbol("head_dim")],
1551                [ElementType::F16],
1552                TensorAccess::ReadWrite,
1553            )?,
1554        ],
1555        outputs: vec![contiguous_tensor_with_alias(
1556            token_hidden_dimensions(),
1557            [ElementType::F16],
1558            TensorAccess::Write,
1559            AliasPolicy::MayAlias { tensor_index: 0 },
1560        )?],
1561        attributes: AttributeSchema::new(BTreeMap::from([
1562            unsigned_attribute("hidden_size")?,
1563            unsigned_attribute("query_heads")?,
1564            unsigned_attribute("kv_heads")?,
1565            unsigned_attribute("head_dim")?,
1566            unsigned_attribute("query_features")?,
1567            unsigned_attribute("kv_features")?,
1568            unsigned_attribute("rope_dim")?,
1569            unsigned_attribute("maximum_context_tokens")?,
1570            positive_rational_attribute("rope_theta")?,
1571            positive_rational_attribute("yarn_factor")?,
1572            unsigned_attribute("yarn_original_context_tokens")?,
1573            positive_rational_attribute("yarn_beta_fast")?,
1574            positive_rational_attribute("yarn_beta_slow")?,
1575            false_bool_attribute("yarn_truncate")?,
1576            nonnegative_unsigned_attribute("sliding_window_tokens")?,
1577            true_bool_attribute("causal")?,
1578            positive_epsilon_attribute("epsilon")?,
1579            nonnegative_unsigned_attribute("layer_index")?,
1580        ]))?,
1581        resources: causal_attention_resources(),
1582        oracle: f16_reference_tolerance()?,
1583        provider: provider_requirement(
1584            GPT_OSS_CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID,
1585            ContractVersion::new(1, 0),
1586        )?,
1587        profile_phase: ProfilePhase::Forward,
1588    };
1589    descriptor.validate()?;
1590    Ok(StandardOperationContract { descriptor })
1591}
1592
1593fn contiguous_tensor(
1594    dimensions: Vec<DimensionConstraint>,
1595    element_types: impl IntoIterator<Item = ElementType>,
1596    access: TensorAccess,
1597) -> Result<TensorContract, VNextError> {
1598    contiguous_tensor_with_alias(dimensions, element_types, access, AliasPolicy::NoAlias)
1599}
1600
1601fn contiguous_tensor_with_alias(
1602    dimensions: Vec<DimensionConstraint>,
1603    element_types: impl IntoIterator<Item = ElementType>,
1604    access: TensorAccess,
1605    alias: AliasPolicy,
1606) -> Result<TensorContract, VNextError> {
1607    TensorContract::new(
1608        dimensions,
1609        element_types.into_iter().collect(),
1610        vec![LayoutConstraint::Contiguous],
1611        access,
1612        alias,
1613    )
1614}
1615
1616fn token_hidden_dimensions() -> Vec<DimensionConstraint> {
1617    vec![
1618        DimensionConstraint::Symbol("tokens".to_owned()),
1619        DimensionConstraint::Symbol("hidden_size".to_owned()),
1620    ]
1621}
1622
1623fn packed_gate_up_dimensions() -> Vec<DimensionConstraint> {
1624    vec![
1625        DimensionConstraint::Exact(2),
1626        DimensionConstraint::Symbol("intermediate_size".to_owned()),
1627        DimensionConstraint::Symbol("hidden_size".to_owned()),
1628    ]
1629}
1630
1631fn intermediate_hidden_dimensions() -> Vec<DimensionConstraint> {
1632    vec![
1633        DimensionConstraint::Symbol("intermediate_size".to_owned()),
1634        DimensionConstraint::Symbol("hidden_size".to_owned()),
1635    ]
1636}
1637
1638fn hidden_intermediate_dimensions() -> Vec<DimensionConstraint> {
1639    vec![
1640        DimensionConstraint::Symbol("hidden_size".to_owned()),
1641        DimensionConstraint::Symbol("intermediate_size".to_owned()),
1642    ]
1643}
1644
1645fn routed_expert_gate_up_dimensions() -> Vec<DimensionConstraint> {
1646    vec![
1647        symbol("expert_count"),
1648        exact(2),
1649        symbol("routed_intermediate_size"),
1650        symbol("hidden_size"),
1651    ]
1652}
1653
1654fn routed_expert_down_dimensions() -> Vec<DimensionConstraint> {
1655    vec![
1656        symbol("expert_count"),
1657        symbol("hidden_size"),
1658        symbol("routed_intermediate_size"),
1659    ]
1660}
1661
1662fn shared_expert_gate_up_dimensions() -> Vec<DimensionConstraint> {
1663    vec![
1664        exact(2),
1665        symbol("shared_intermediate_size"),
1666        symbol("hidden_size"),
1667    ]
1668}
1669
1670fn shared_expert_down_dimensions() -> Vec<DimensionConstraint> {
1671    vec![symbol("hidden_size"), symbol("shared_intermediate_size")]
1672}
1673
1674fn no_auxiliary_resources() -> ResourceRequirements {
1675    ResourceRequirements {
1676        minimum_value_alignment_bytes: 16,
1677        scratch: ResourcePresenceRequirement::Forbidden,
1678        binding: ResourcePresenceRequirement::Forbidden,
1679        persistent: ResourcePresenceRequirement::Forbidden,
1680    }
1681}
1682
1683fn attention_resources() -> ResourceRequirements {
1684    ResourceRequirements {
1685        minimum_value_alignment_bytes: 16,
1686        scratch: ResourcePresenceRequirement::Required,
1687        binding: ResourcePresenceRequirement::Optional,
1688        persistent: ResourcePresenceRequirement::Forbidden,
1689    }
1690}
1691
1692fn causal_attention_resources() -> ResourceRequirements {
1693    ResourceRequirements {
1694        minimum_value_alignment_bytes: 16,
1695        scratch: ResourcePresenceRequirement::Required,
1696        binding: ResourcePresenceRequirement::Required,
1697        persistent: ResourcePresenceRequirement::Forbidden,
1698    }
1699}
1700
1701fn symbol(name: &str) -> DimensionConstraint {
1702    DimensionConstraint::Symbol(name.to_owned())
1703}
1704
1705const fn exact(value: u64) -> DimensionConstraint {
1706    DimensionConstraint::Exact(value)
1707}
1708
1709fn provider_requirement(
1710    capability: &str,
1711    minimum_version: ContractVersion,
1712) -> Result<ProviderRequirement, VNextError> {
1713    Ok(ProviderRequirement {
1714        minimum_version,
1715        required_capabilities: BTreeSet::from([CapabilityId::new(capability)?]),
1716    })
1717}
1718
1719fn f16_reference_tolerance() -> Result<OracleSpec, VNextError> {
1720    Ok(OracleSpec::RelativeTolerance {
1721        tolerance: CanonicalRational::new(1, 1_000)?,
1722    })
1723}
1724
1725fn f32_reference_tolerance() -> Result<OracleSpec, VNextError> {
1726    Ok(OracleSpec::RelativeTolerance {
1727        tolerance: CanonicalRational::new(1, 100_000)?,
1728    })
1729}
1730
1731fn unsigned_attribute(name: &str) -> Result<(AttributeId, AttributeSpec), VNextError> {
1732    Ok((
1733        AttributeId::new(name)?,
1734        AttributeSpec {
1735            value_kind: AttributeValueKind::Unsigned,
1736            required: true,
1737            constraint: AttributeConstraint::UnsignedRange {
1738                minimum: 1,
1739                maximum: u32::MAX as u64,
1740            },
1741        },
1742    ))
1743}
1744
1745fn nonnegative_unsigned_attribute(name: &str) -> Result<(AttributeId, AttributeSpec), VNextError> {
1746    Ok((
1747        AttributeId::new(name)?,
1748        AttributeSpec {
1749            value_kind: AttributeValueKind::Unsigned,
1750            required: true,
1751            constraint: AttributeConstraint::UnsignedRange {
1752                minimum: 0,
1753                maximum: u32::MAX as u64,
1754            },
1755        },
1756    ))
1757}
1758
1759fn unconstrained_bool_attribute(name: &str) -> Result<(AttributeId, AttributeSpec), VNextError> {
1760    Ok((
1761        AttributeId::new(name)?,
1762        AttributeSpec {
1763            value_kind: AttributeValueKind::Bool,
1764            required: true,
1765            constraint: AttributeConstraint::None,
1766        },
1767    ))
1768}
1769
1770fn true_bool_attribute(name: &str) -> Result<(AttributeId, AttributeSpec), VNextError> {
1771    Ok((
1772        AttributeId::new(name)?,
1773        AttributeSpec {
1774            value_kind: AttributeValueKind::Bool,
1775            required: true,
1776            constraint: AttributeConstraint::BoolEquals(true),
1777        },
1778    ))
1779}
1780
1781fn false_bool_attribute(name: &str) -> Result<(AttributeId, AttributeSpec), VNextError> {
1782    Ok((
1783        AttributeId::new(name)?,
1784        AttributeSpec {
1785            value_kind: AttributeValueKind::Bool,
1786            required: true,
1787            constraint: AttributeConstraint::BoolEquals(false),
1788        },
1789    ))
1790}
1791
1792fn exact_rational_attribute(
1793    name: &str,
1794    numerator: i64,
1795    denominator: u64,
1796) -> Result<(AttributeId, AttributeSpec), VNextError> {
1797    let value = CanonicalRational::new(numerator, denominator)?;
1798    Ok((
1799        AttributeId::new(name)?,
1800        AttributeSpec {
1801            value_kind: AttributeValueKind::Rational,
1802            required: true,
1803            constraint: AttributeConstraint::RationalRange {
1804                minimum: value,
1805                maximum: value,
1806            },
1807        },
1808    ))
1809}
1810
1811fn positive_rational_attribute(name: &str) -> Result<(AttributeId, AttributeSpec), VNextError> {
1812    Ok((
1813        AttributeId::new(name)?,
1814        AttributeSpec {
1815            value_kind: AttributeValueKind::Rational,
1816            required: true,
1817            constraint: AttributeConstraint::RationalRange {
1818                minimum: CanonicalRational::new(1, u64::MAX)?,
1819                maximum: CanonicalRational::new(i64::MAX, 1)?,
1820            },
1821        },
1822    ))
1823}
1824
1825fn positive_epsilon_attribute(name: &str) -> Result<(AttributeId, AttributeSpec), VNextError> {
1826    Ok((
1827        AttributeId::new(name)?,
1828        AttributeSpec {
1829            value_kind: AttributeValueKind::Rational,
1830            required: true,
1831            constraint: AttributeConstraint::RationalRange {
1832                minimum: CanonicalRational::new(1, 1_000_000_000_000)?,
1833                maximum: CanonicalRational::new(1, 1)?,
1834            },
1835        },
1836    ))
1837}
1838
1839fn text_choices_attribute(
1840    name: &str,
1841    values: impl IntoIterator<Item = &'static str>,
1842) -> Result<(AttributeId, AttributeSpec), VNextError> {
1843    Ok((
1844        AttributeId::new(name)?,
1845        AttributeSpec {
1846            value_kind: AttributeValueKind::Text,
1847            required: true,
1848            constraint: AttributeConstraint::TextChoices {
1849                values: values.into_iter().map(str::to_owned).collect(),
1850            },
1851        },
1852    ))
1853}
1854
1855#[cfg(test)]
1856mod tests {
1857    use super::*;
1858
1859    #[test]
1860    fn gated_delta_recurrent_only_capability_never_claims_chunked_scan() {
1861        let capabilities = GatedDeltaExecutionCapabilities::recurrent_only();
1862        assert_eq!(
1863            capabilities
1864                .select(1, GatedDeltaExecutionPreference::ChunkedScan)
1865                .unwrap(),
1866            GatedDeltaExecutionForm::RecurrentScan
1867        );
1868        assert_eq!(
1869            capabilities
1870                .select(64, GatedDeltaExecutionPreference::ChunkedScan)
1871                .unwrap(),
1872            GatedDeltaExecutionForm::RecurrentScan
1873        );
1874    }
1875
1876    #[test]
1877    fn gated_delta_chunk_plan_preserves_exact_tail_boundaries() {
1878        let capabilities = GatedDeltaExecutionCapabilities::with_chunked_scan(64).unwrap();
1879        assert_eq!(
1880            capabilities
1881                .select(1, GatedDeltaExecutionPreference::ChunkedScan)
1882                .unwrap(),
1883            GatedDeltaExecutionForm::RecurrentScan
1884        );
1885        assert_eq!(
1886            capabilities
1887                .select(64, GatedDeltaExecutionPreference::RecurrentScan)
1888                .unwrap(),
1889            GatedDeltaExecutionForm::RecurrentScan
1890        );
1891        for (tokens, chunks, final_tokens) in [(2, 1, 2), (64, 1, 64), (65, 2, 1)] {
1892            let GatedDeltaExecutionForm::ChunkedScan(plan) = capabilities
1893                .select(tokens, GatedDeltaExecutionPreference::ChunkedScan)
1894                .unwrap()
1895            else {
1896                panic!("{tokens} tokens must select chunked scan");
1897            };
1898            assert_eq!(plan.token_count(), tokens);
1899            assert_eq!(plan.chunk_size(), 64);
1900            assert_eq!(plan.chunk_count(), chunks);
1901            assert_eq!(plan.final_chunk_tokens(), final_tokens);
1902            assert_eq!(
1903                GatedDeltaExecutionForm::ChunkedScan(plan).as_str(),
1904                "chunked_scan"
1905            );
1906        }
1907    }
1908
1909    #[test]
1910    fn gated_delta_execution_capabilities_reject_invalid_domains() {
1911        assert!(GatedDeltaExecutionCapabilities::with_chunked_scan(0).is_err());
1912        assert!(GatedDeltaExecutionCapabilities::recurrent_only()
1913            .select(0, GatedDeltaExecutionPreference::RecurrentScan)
1914            .is_err());
1915    }
1916
1917    #[test]
1918    fn token_embedding_contract_is_backend_and_model_neutral() {
1919        let contract = token_embedding_contract().unwrap();
1920        let descriptor = contract.descriptor();
1921        assert_eq!(descriptor.id.as_str(), TOKEN_EMBEDDING_OPERATION_ID);
1922        assert_eq!(descriptor.fingerprint().unwrap().len(), 64);
1923        contract
1924            .validate_signature(&descriptor.inputs, &descriptor.outputs)
1925            .unwrap();
1926    }
1927
1928    #[test]
1929    fn int8_kv_contracts_require_both_typed_states_without_changing_other_boundaries() {
1930        for (f16, quantized) in [
1931            (
1932                causal_paged_attention_contract().unwrap(),
1933                causal_paged_attention_int8_kv_contract().unwrap(),
1934            ),
1935            (
1936                causal_paged_attention_f32_master_contract().unwrap(),
1937                causal_paged_attention_f32_master_int8_kv_contract().unwrap(),
1938            ),
1939        ] {
1940            let old = f16.descriptor();
1941            let new = quantized.descriptor();
1942            assert_eq!(&old.inputs[..8], &new.inputs[..8]);
1943            assert_eq!(old.outputs, new.outputs);
1944            assert_eq!(old.attributes, new.attributes);
1945            assert_eq!(new.inputs.len(), 10);
1946            assert_eq!(
1947                new.inputs[8].element_types(),
1948                &BTreeSet::from([ElementType::I8])
1949            );
1950            assert_eq!(
1951                new.inputs[9].element_types(),
1952                &BTreeSet::from([ElementType::F32])
1953            );
1954            assert_eq!(new.inputs[8].access(), TensorAccess::ReadWrite);
1955            assert_eq!(new.inputs[9].access(), TensorAccess::ReadWrite);
1956            assert_ne!(old.fingerprint().unwrap(), new.fingerprint().unwrap());
1957            quantized
1958                .validate_signature(&new.inputs, &new.outputs)
1959                .unwrap();
1960            assert!(quantized
1961                .validate_signature(&old.inputs, &old.outputs)
1962                .is_err());
1963        }
1964    }
1965
1966    #[test]
1967    fn fp32_master_contracts_form_one_exact_mixed_precision_chain() {
1968        let contracts = [
1969            token_embedding_f32_master_contract().unwrap(),
1970            gated_delta_recurrent_attention_f32_master_contract().unwrap(),
1971            causal_paged_attention_f32_master_contract().unwrap(),
1972            rms_norm_f32_to_f16_contract().unwrap(),
1973            residual_add_f32_f16_contract().unwrap(),
1974            rms_norm_f32_contract().unwrap(),
1975            last_token_dense_linear_f32_contract().unwrap(),
1976            last_token_masked_argmax_f32_contract().unwrap(),
1977        ];
1978        let ids = contracts
1979            .iter()
1980            .map(|contract| contract.descriptor().id.as_str())
1981            .collect::<BTreeSet<_>>();
1982        assert_eq!(ids.len(), contracts.len());
1983        for contract in &contracts {
1984            let descriptor = contract.descriptor();
1985            let minor =
1986                u16::from(descriptor.id.as_str() == TOKEN_EMBEDDING_F32_MASTER_OPERATION_ID);
1987            assert_eq!(descriptor.version, ContractVersion::new(1, minor));
1988            assert_eq!(descriptor.provider.required_capabilities.len(), 1);
1989            contract
1990                .validate_signature(&descriptor.inputs, &descriptor.outputs)
1991                .unwrap();
1992        }
1993
1994        let embedding = contracts[0].descriptor();
1995        assert_eq!(
1996            embedding.outputs[0].element_types(),
1997            &BTreeSet::from([ElementType::F32])
1998        );
1999        for attention in [&contracts[1], &contracts[2]] {
2000            let descriptor = attention.descriptor();
2001            assert_eq!(
2002                descriptor.inputs[0].element_types(),
2003                &BTreeSet::from([ElementType::F32])
2004            );
2005            assert_eq!(
2006                descriptor.outputs[0].element_types(),
2007                &BTreeSet::from([ElementType::F32])
2008            );
2009            assert_eq!(
2010                descriptor.outputs[0].alias(),
2011                &AliasPolicy::MayAlias { tensor_index: 0 }
2012            );
2013        }
2014
2015        let branch_norm = contracts[3].descriptor();
2016        assert_eq!(
2017            branch_norm.inputs[0].element_types(),
2018            &BTreeSet::from([ElementType::F32])
2019        );
2020        assert_eq!(
2021            branch_norm.outputs[0].element_types(),
2022            &BTreeSet::from([ElementType::F16])
2023        );
2024        let residual = contracts[4].descriptor();
2025        assert_eq!(
2026            residual.inputs[0].element_types(),
2027            &BTreeSet::from([ElementType::F32])
2028        );
2029        assert_eq!(
2030            residual.inputs[1].element_types(),
2031            &BTreeSet::from([ElementType::F16])
2032        );
2033        assert_eq!(
2034            residual.outputs[0].element_types(),
2035            &BTreeSet::from([ElementType::F32])
2036        );
2037
2038        let final_norm = contracts[5].descriptor();
2039        let head = contracts[6].descriptor();
2040        let argmax = contracts[7].descriptor();
2041        for tensor in [
2042            &final_norm.inputs[0],
2043            &final_norm.outputs[0],
2044            &head.inputs[0],
2045            &head.outputs[0],
2046            &argmax.inputs[0],
2047        ] {
2048            assert_eq!(tensor.element_types(), &BTreeSet::from([ElementType::F32]));
2049        }
2050
2051        assert!(!ids.contains(TOKEN_EMBEDDING_OPERATION_ID));
2052        assert!(!ids.contains(RMS_NORM_OPERATION_ID));
2053        assert!(!ids.contains(RESIDUAL_ADD_OPERATION_ID));
2054    }
2055
2056    #[test]
2057    fn last_token_dense_linear_contract_is_backend_and_model_neutral() {
2058        let contract = last_token_dense_linear_contract().unwrap();
2059        let descriptor = contract.descriptor();
2060        assert_eq!(descriptor.id.as_str(), LAST_TOKEN_DENSE_LINEAR_OPERATION_ID);
2061        assert_eq!(descriptor.version, ContractVersion::new(1, 1));
2062        assert_eq!(
2063            descriptor.resources.scratch,
2064            ResourcePresenceRequirement::Optional
2065        );
2066        assert_eq!(
2067            descriptor.outputs[0].dimensions(),
2068            &[
2069                DimensionConstraint::Exact(1),
2070                DimensionConstraint::Symbol("out_features".to_owned()),
2071            ]
2072        );
2073        contract
2074            .validate_signature(&descriptor.inputs, &descriptor.outputs)
2075            .unwrap();
2076    }
2077
2078    #[test]
2079    fn last_token_masked_argmax_contract_keeps_policy_in_typed_inputs() {
2080        let contract = last_token_masked_argmax_contract().unwrap();
2081        let descriptor = contract.descriptor();
2082        assert_eq!(
2083            descriptor.id.as_str(),
2084            LAST_TOKEN_MASKED_ARGMAX_OPERATION_ID
2085        );
2086        assert_eq!(descriptor.version, ContractVersion::new(3, 0));
2087        assert_eq!(descriptor.inputs.len(), 5);
2088        assert_eq!(descriptor.inputs[0].access(), TensorAccess::Read);
2089        assert_eq!(
2090            descriptor.resources.scratch,
2091            ResourcePresenceRequirement::Required
2092        );
2093        assert_eq!(
2094            descriptor.resources.binding,
2095            ResourcePresenceRequirement::Forbidden
2096        );
2097        assert_eq!(
2098            descriptor.resources.persistent,
2099            ResourcePresenceRequirement::Forbidden
2100        );
2101        assert_eq!(
2102            descriptor.inputs[1].element_types(),
2103            &BTreeSet::from([ElementType::U8])
2104        );
2105        assert_eq!(
2106            descriptor.inputs[2].element_types(),
2107            &BTreeSet::from([ElementType::U32])
2108        );
2109        assert_eq!(
2110            descriptor.inputs[3].dimensions(),
2111            &[DimensionConstraint::Exact(2)]
2112        );
2113        assert_eq!(
2114            descriptor.inputs[4].element_types(),
2115            &BTreeSet::from([ElementType::F32])
2116        );
2117        assert_eq!(
2118            descriptor.outputs[0].element_types(),
2119            &BTreeSet::from([ElementType::U32])
2120        );
2121        assert_eq!(
2122            descriptor.outputs[0].dimensions(),
2123            &[DimensionConstraint::Exact(1)]
2124        );
2125        contract
2126            .validate_signature(&descriptor.inputs, &descriptor.outputs)
2127            .unwrap();
2128    }
2129
2130    #[test]
2131    fn transformer_primitives_have_explicit_math_and_resource_boundaries() {
2132        let contracts = [
2133            rms_norm_contract().unwrap(),
2134            dense_linear_contract().unwrap(),
2135            dense_swiglu_contract().unwrap(),
2136            residual_add_contract().unwrap(),
2137        ];
2138        for contract in &contracts {
2139            let descriptor = contract.descriptor();
2140            assert_eq!(descriptor.fingerprint().unwrap().len(), 64);
2141            contract
2142                .validate_signature(&descriptor.inputs, &descriptor.outputs)
2143                .unwrap();
2144        }
2145        assert_eq!(
2146            contracts[2].descriptor().resources.scratch,
2147            ResourcePresenceRequirement::Required
2148        );
2149        assert_eq!(
2150            contracts[3].descriptor().outputs[0].alias(),
2151            &AliasPolicy::MayAlias { tensor_index: 0 }
2152        );
2153    }
2154
2155    #[test]
2156    fn hybrid_vnorm_simple_ops_have_typed_shapes_attributes_and_aliasing() {
2157        let geglu = dense_geglu_tanh_contract().unwrap();
2158        let geglu_descriptor = geglu.descriptor();
2159        assert_eq!(geglu_descriptor.id.as_str(), DENSE_GEGLU_TANH_OPERATION_ID);
2160        assert_eq!(geglu_descriptor.version, ContractVersion::new(1, 0));
2161        assert_eq!(geglu_descriptor.inputs.len(), 4);
2162        assert_eq!(
2163            geglu_descriptor.inputs[0].dimensions(),
2164            &[symbol("tokens"), symbol("hidden_size")]
2165        );
2166        for projection in [&geglu_descriptor.inputs[1], &geglu_descriptor.inputs[2]] {
2167            assert_eq!(
2168                projection.dimensions(),
2169                &[symbol("intermediate_size"), symbol("hidden_size")]
2170            );
2171        }
2172        assert_eq!(
2173            geglu_descriptor.inputs[3].dimensions(),
2174            &[symbol("hidden_size"), symbol("intermediate_size")]
2175        );
2176        assert_eq!(
2177            geglu_descriptor.outputs[0].dimensions(),
2178            &[symbol("tokens"), symbol("hidden_size")]
2179        );
2180        assert_eq!(
2181            geglu_descriptor.outputs[0].element_types(),
2182            &BTreeSet::from([ElementType::F16])
2183        );
2184        assert_eq!(
2185            geglu_descriptor.resources.scratch,
2186            ResourcePresenceRequirement::Required
2187        );
2188        assert_eq!(
2189            geglu_descriptor.provider.required_capabilities,
2190            BTreeSet::from([CapabilityId::new(DENSE_GEGLU_TANH_F16_CAPABILITY_ID).unwrap()])
2191        );
2192
2193        let scale = constant_scale_contract().unwrap();
2194        let scale_descriptor = scale.descriptor();
2195        assert_eq!(scale_descriptor.id.as_str(), CONSTANT_SCALE_OPERATION_ID);
2196        assert_eq!(scale_descriptor.inputs.len(), 1);
2197        assert_eq!(
2198            scale_descriptor.outputs[0].alias(),
2199            &AliasPolicy::MustAlias { tensor_index: 0 }
2200        );
2201        assert!(scale_descriptor
2202            .attributes
2203            .entries()
2204            .contains_key(&AttributeId::new("scale").unwrap()));
2205        assert_eq!(scale_descriptor.resources, no_auxiliary_resources());
2206
2207        let softcap = logit_softcap_contract().unwrap();
2208        let softcap_descriptor = softcap.descriptor();
2209        assert_eq!(softcap_descriptor.id.as_str(), LOGIT_SOFTCAP_OPERATION_ID);
2210        assert_eq!(
2211            softcap_descriptor.inputs[0].dimensions(),
2212            &[exact(1), symbol("vocab_size")]
2213        );
2214        assert_eq!(
2215            softcap_descriptor.outputs[0].alias(),
2216            &AliasPolicy::MustAlias { tensor_index: 0 }
2217        );
2218        assert!(softcap_descriptor
2219            .attributes
2220            .entries()
2221            .contains_key(&AttributeId::new("cap").unwrap()));
2222        assert_eq!(softcap_descriptor.resources, no_auxiliary_resources());
2223
2224        for (descriptor, attribute) in [(scale_descriptor, "scale"), (softcap_descriptor, "cap")] {
2225            let AttributeConstraint::RationalRange { minimum, maximum } = &descriptor
2226                .attributes
2227                .entries()
2228                .get(&AttributeId::new(attribute).unwrap())
2229                .unwrap()
2230                .constraint
2231            else {
2232                panic!("{attribute} must have a typed rational range");
2233            };
2234            assert!(minimum.numerator() > 0);
2235            assert!(maximum >= minimum);
2236        }
2237
2238        for contract in [&geglu, &scale, &softcap] {
2239            let descriptor = contract.descriptor();
2240            assert_eq!(descriptor.fingerprint().unwrap().len(), 64);
2241            contract
2242                .validate_signature(&descriptor.inputs, &descriptor.outputs)
2243                .unwrap();
2244            assert!(contract
2245                .validate_signature(&[], &descriptor.outputs)
2246                .is_err());
2247        }
2248    }
2249
2250    #[test]
2251    fn routed_shared_moe_contract_keeps_fusion_and_weight_abi_generic() {
2252        let contract = routed_shared_swiglu_moe_contract().unwrap();
2253        let descriptor = contract.descriptor();
2254
2255        assert_eq!(
2256            descriptor.id.as_str(),
2257            ROUTED_SHARED_SWIGLU_MOE_OPERATION_ID
2258        );
2259        assert_eq!(descriptor.version, ContractVersion::new(1, 0));
2260        assert_eq!(descriptor.inputs.len(), 7);
2261        assert_eq!(
2262            descriptor.inputs[2].dimensions(),
2263            &[
2264                symbol("expert_count"),
2265                exact(2),
2266                symbol("routed_intermediate_size"),
2267                symbol("hidden_size"),
2268            ]
2269        );
2270        assert_eq!(
2271            descriptor.inputs[3].dimensions(),
2272            &[
2273                symbol("expert_count"),
2274                symbol("hidden_size"),
2275                symbol("routed_intermediate_size"),
2276            ]
2277        );
2278        assert_eq!(
2279            descriptor.inputs[4].dimensions(),
2280            &[exact(1), symbol("hidden_size")]
2281        );
2282        assert_eq!(
2283            descriptor.resources.scratch,
2284            ResourcePresenceRequirement::Required
2285        );
2286        assert_eq!(
2287            descriptor.resources.persistent,
2288            ResourcePresenceRequirement::Forbidden
2289        );
2290        assert_eq!(
2291            descriptor.provider.required_capabilities,
2292            BTreeSet::from([
2293                CapabilityId::new(ROUTED_SHARED_SWIGLU_MOE_F16_CAPABILITY_ID).unwrap()
2294            ])
2295        );
2296        for attribute in [
2297            "hidden_size",
2298            "expert_count",
2299            "experts_per_token",
2300            "routed_intermediate_size",
2301            "shared_intermediate_size",
2302            "normalize_topk",
2303        ] {
2304            assert!(
2305                descriptor
2306                    .attributes
2307                    .entries()
2308                    .contains_key(&AttributeId::new(attribute).unwrap()),
2309                "missing typed MoE attribute {attribute}"
2310            );
2311        }
2312        assert_eq!(descriptor.fingerprint().unwrap().len(), 64);
2313        contract
2314            .validate_signature(&descriptor.inputs, &descriptor.outputs)
2315            .unwrap();
2316    }
2317
2318    #[test]
2319    fn routed_only_moe_contract_has_no_shared_expert_abi() {
2320        let contract = routed_swiglu_moe_contract().unwrap();
2321        let descriptor = contract.descriptor();
2322
2323        assert_eq!(descriptor.id.as_str(), ROUTED_SWIGLU_MOE_OPERATION_ID);
2324        assert_eq!(descriptor.version, ContractVersion::new(1, 0));
2325        assert_eq!(descriptor.inputs.len(), 4);
2326        assert_eq!(
2327            descriptor.inputs[2].dimensions(),
2328            &[
2329                symbol("expert_count"),
2330                exact(2),
2331                symbol("routed_intermediate_size"),
2332                symbol("hidden_size"),
2333            ]
2334        );
2335        assert_eq!(
2336            descriptor.inputs[3].dimensions(),
2337            &[
2338                symbol("expert_count"),
2339                symbol("hidden_size"),
2340                symbol("routed_intermediate_size"),
2341            ]
2342        );
2343        assert_eq!(
2344            descriptor.resources.scratch,
2345            ResourcePresenceRequirement::Required
2346        );
2347        assert_eq!(
2348            descriptor.provider.required_capabilities,
2349            BTreeSet::from([CapabilityId::new(ROUTED_SWIGLU_MOE_F16_CAPABILITY_ID).unwrap()])
2350        );
2351        for attribute in [
2352            "hidden_size",
2353            "expert_count",
2354            "experts_per_token",
2355            "routed_intermediate_size",
2356            "normalize_topk",
2357        ] {
2358            assert!(
2359                descriptor
2360                    .attributes
2361                    .entries()
2362                    .contains_key(&AttributeId::new(attribute).unwrap()),
2363                "missing typed routed-only MoE attribute {attribute}"
2364            );
2365        }
2366        assert!(!descriptor
2367            .attributes
2368            .entries()
2369            .contains_key(&AttributeId::new("shared_intermediate_size").unwrap()));
2370        assert_eq!(descriptor.fingerprint().unwrap().len(), 64);
2371        contract
2372            .validate_signature(&descriptor.inputs, &descriptor.outputs)
2373            .unwrap();
2374    }
2375
2376    #[test]
2377    fn gpt_oss_routed_clamped_swiglu_contract_has_exact_logical_bf16_abi() {
2378        let contract = gpt_oss_routed_clamped_swiglu_moe_contract().unwrap();
2379        let descriptor = contract.descriptor();
2380
2381        assert_eq!(
2382            descriptor.id.as_str(),
2383            GPT_OSS_ROUTED_CLAMPED_SWIGLU_MOE_OPERATION_ID
2384        );
2385        assert_eq!(descriptor.version, ContractVersion::new(1, 0));
2386        assert_eq!(descriptor.inputs.len(), 7);
2387        assert_eq!(
2388            descriptor.inputs[0].dimensions(),
2389            &[symbol("tokens"), symbol("hidden_size")]
2390        );
2391        assert_eq!(
2392            descriptor.inputs[1].dimensions(),
2393            &[symbol("expert_count"), symbol("hidden_size")]
2394        );
2395        assert_eq!(descriptor.inputs[2].dimensions(), &[symbol("expert_count")]);
2396        assert_eq!(
2397            descriptor.inputs[3].dimensions(),
2398            &[
2399                symbol("expert_count"),
2400                symbol("gate_up_features"),
2401                symbol("hidden_size"),
2402            ]
2403        );
2404        assert_eq!(
2405            descriptor.inputs[4].dimensions(),
2406            &[symbol("expert_count"), symbol("gate_up_features")]
2407        );
2408        assert_eq!(
2409            descriptor.inputs[5].dimensions(),
2410            &[
2411                symbol("expert_count"),
2412                symbol("hidden_size"),
2413                symbol("intermediate_size"),
2414            ]
2415        );
2416        assert_eq!(
2417            descriptor.inputs[6].dimensions(),
2418            &[symbol("expert_count"), symbol("hidden_size")]
2419        );
2420        assert_eq!(
2421            descriptor.inputs[0].element_types(),
2422            &BTreeSet::from([ElementType::F16])
2423        );
2424        for input in &descriptor.inputs[1..] {
2425            assert_eq!(input.element_types(), &BTreeSet::from([ElementType::Bf16]));
2426        }
2427        assert!(descriptor
2428            .inputs
2429            .iter()
2430            .all(|input| input.access() == TensorAccess::Read));
2431        assert_eq!(
2432            descriptor.outputs[0].dimensions(),
2433            &[symbol("tokens"), symbol("hidden_size")]
2434        );
2435        assert_eq!(
2436            descriptor.outputs[0].element_types(),
2437            &BTreeSet::from([ElementType::F16])
2438        );
2439        assert_eq!(descriptor.outputs[0].alias(), &AliasPolicy::NoAlias);
2440
2441        assert_eq!(
2442            descriptor
2443                .attributes
2444                .entries()
2445                .keys()
2446                .map(AttributeId::as_str)
2447                .collect::<BTreeSet<_>>(),
2448            BTreeSet::from([
2449                "down_bias_before_route_reduction",
2450                "expert_count",
2451                "experts_per_token",
2452                "gate_up_features",
2453                "gate_up_interleaved",
2454                "hidden_size",
2455                "intermediate_size",
2456                "normalize_topk",
2457                "swiglu_limit",
2458            ])
2459        );
2460        for attribute in [
2461            "normalize_topk",
2462            "gate_up_interleaved",
2463            "down_bias_before_route_reduction",
2464        ] {
2465            assert_eq!(
2466                descriptor
2467                    .attributes
2468                    .entries()
2469                    .get(&AttributeId::new(attribute).unwrap())
2470                    .unwrap()
2471                    .constraint,
2472                AttributeConstraint::BoolEquals(true)
2473            );
2474        }
2475        let seven = CanonicalRational::new(7, 1).unwrap();
2476        assert_eq!(
2477            descriptor
2478                .attributes
2479                .entries()
2480                .get(&AttributeId::new("swiglu_limit").unwrap())
2481                .unwrap()
2482                .constraint,
2483            AttributeConstraint::RationalRange {
2484                minimum: seven,
2485                maximum: seven,
2486            }
2487        );
2488        assert_eq!(
2489            descriptor.resources.scratch,
2490            ResourcePresenceRequirement::Required
2491        );
2492        assert_eq!(
2493            descriptor.resources.binding,
2494            ResourcePresenceRequirement::Forbidden
2495        );
2496        assert_eq!(
2497            descriptor.resources.persistent,
2498            ResourcePresenceRequirement::Forbidden
2499        );
2500        assert_eq!(
2501            descriptor.provider.required_capabilities,
2502            BTreeSet::from([CapabilityId::new(
2503                GPT_OSS_ROUTED_CLAMPED_SWIGLU_MOE_MXFP4_BF16_CAPABILITY_ID
2504            )
2505            .unwrap()])
2506        );
2507        assert_eq!(descriptor.fingerprint().unwrap().len(), 64);
2508        contract
2509            .validate_signature(&descriptor.inputs, &descriptor.outputs)
2510            .unwrap();
2511    }
2512
2513    #[test]
2514    fn gpt_oss_causal_attention_contract_has_exact_bias_sink_and_yarn_abi() {
2515        let contract = gpt_oss_causal_paged_attention_contract().unwrap();
2516        let descriptor = contract.descriptor();
2517
2518        assert_eq!(
2519            descriptor.id.as_str(),
2520            GPT_OSS_CAUSAL_PAGED_ATTENTION_OPERATION_ID
2521        );
2522        assert_eq!(descriptor.version, ContractVersion::new(1, 0));
2523        assert_eq!(descriptor.inputs.len(), 12);
2524        let expected_dimensions = [
2525            vec![symbol("tokens"), symbol("hidden_size")],
2526            vec![symbol("hidden_size")],
2527            vec![symbol("query_features"), symbol("hidden_size")],
2528            vec![symbol("kv_features"), symbol("hidden_size")],
2529            vec![symbol("kv_features"), symbol("hidden_size")],
2530            vec![symbol("hidden_size"), symbol("query_features")],
2531            vec![symbol("query_features")],
2532            vec![symbol("kv_features")],
2533            vec![symbol("kv_features")],
2534            vec![symbol("hidden_size")],
2535            vec![symbol("query_heads")],
2536            vec![exact(2), symbol("kv_heads"), symbol("head_dim")],
2537        ];
2538        for (input, expected) in descriptor.inputs.iter().zip(expected_dimensions) {
2539            assert_eq!(input.dimensions(), expected);
2540            assert_eq!(input.element_types(), &BTreeSet::from([ElementType::F16]));
2541        }
2542        assert!(descriptor.inputs[..11]
2543            .iter()
2544            .all(|input| input.access() == TensorAccess::Read));
2545        assert_eq!(descriptor.inputs[11].access(), TensorAccess::ReadWrite);
2546        assert_eq!(
2547            descriptor.outputs[0].dimensions(),
2548            &[symbol("tokens"), symbol("hidden_size")]
2549        );
2550        assert_eq!(
2551            descriptor.outputs[0].element_types(),
2552            &BTreeSet::from([ElementType::F16])
2553        );
2554        assert_eq!(
2555            descriptor.outputs[0].alias(),
2556            &AliasPolicy::MayAlias { tensor_index: 0 }
2557        );
2558
2559        assert_eq!(
2560            descriptor
2561                .attributes
2562                .entries()
2563                .keys()
2564                .map(AttributeId::as_str)
2565                .collect::<BTreeSet<_>>(),
2566            BTreeSet::from([
2567                "causal",
2568                "epsilon",
2569                "head_dim",
2570                "hidden_size",
2571                "kv_features",
2572                "kv_heads",
2573                "layer_index",
2574                "maximum_context_tokens",
2575                "query_features",
2576                "query_heads",
2577                "rope_dim",
2578                "rope_theta",
2579                "sliding_window_tokens",
2580                "yarn_beta_fast",
2581                "yarn_beta_slow",
2582                "yarn_factor",
2583                "yarn_original_context_tokens",
2584                "yarn_truncate",
2585            ])
2586        );
2587        assert_eq!(
2588            descriptor
2589                .attributes
2590                .entries()
2591                .get(&AttributeId::new("yarn_truncate").unwrap())
2592                .unwrap()
2593                .constraint,
2594            AttributeConstraint::BoolEquals(false)
2595        );
2596        assert_eq!(
2597            descriptor
2598                .attributes
2599                .entries()
2600                .get(&AttributeId::new("causal").unwrap())
2601                .unwrap()
2602                .constraint,
2603            AttributeConstraint::BoolEquals(true)
2604        );
2605        assert_eq!(
2606            descriptor
2607                .attributes
2608                .entries()
2609                .get(&AttributeId::new("sliding_window_tokens").unwrap())
2610                .unwrap()
2611                .constraint,
2612            AttributeConstraint::UnsignedRange {
2613                minimum: 0,
2614                maximum: u32::MAX as u64,
2615            }
2616        );
2617        assert_eq!(
2618            descriptor.resources.scratch,
2619            ResourcePresenceRequirement::Required
2620        );
2621        assert_eq!(
2622            descriptor.resources.binding,
2623            ResourcePresenceRequirement::Required
2624        );
2625        assert_eq!(
2626            descriptor.resources.persistent,
2627            ResourcePresenceRequirement::Forbidden
2628        );
2629        assert_eq!(
2630            descriptor.provider.required_capabilities,
2631            BTreeSet::from([
2632                CapabilityId::new(GPT_OSS_CAUSAL_PAGED_ATTENTION_F16_CAPABILITY_ID).unwrap()
2633            ])
2634        );
2635        assert_eq!(descriptor.fingerprint().unwrap().len(), 64);
2636        contract
2637            .validate_signature(&descriptor.inputs, &descriptor.outputs)
2638            .unwrap();
2639    }
2640
2641    #[test]
2642    fn attention_contracts_fix_weight_order_state_mutability_and_scratch() {
2643        let linear = gated_delta_recurrent_attention_contract().unwrap();
2644        let full = causal_paged_attention_contract().unwrap();
2645        for contract in [&linear, &full] {
2646            let descriptor = contract.descriptor();
2647            assert_eq!(
2648                descriptor.resources.scratch,
2649                ResourcePresenceRequirement::Required
2650            );
2651            assert_eq!(
2652                descriptor.outputs[0].alias(),
2653                &AliasPolicy::MayAlias { tensor_index: 0 }
2654            );
2655            assert_eq!(descriptor.fingerprint().unwrap().len(), 64);
2656            contract
2657                .validate_signature(&descriptor.inputs, &descriptor.outputs)
2658                .unwrap();
2659        }
2660        assert_eq!(linear.descriptor().inputs.len(), 10);
2661        assert_eq!(linear.descriptor().version, ContractVersion::new(6, 0));
2662        assert_eq!(
2663            linear.descriptor().resources.binding,
2664            ResourcePresenceRequirement::Optional
2665        );
2666        assert_eq!(
2667            full.descriptor().resources.binding,
2668            ResourcePresenceRequirement::Required
2669        );
2670        assert_eq!(
2671            linear.descriptor().provider.minimum_version,
2672            ContractVersion::new(6, 0)
2673        );
2674        for (name, values) in [
2675            (
2676                "decay_parameterization",
2677                GatedDeltaDecayParameterization::ALL
2678                    .map(|value| value.as_str().to_owned())
2679                    .into_iter()
2680                    .collect(),
2681            ),
2682            (
2683                "value_head_mapping",
2684                GatedDeltaValueHeadMapping::ALL
2685                    .map(|value| value.as_str().to_owned())
2686                    .into_iter()
2687                    .collect(),
2688            ),
2689        ] {
2690            assert_eq!(
2691                linear
2692                    .descriptor()
2693                    .attributes
2694                    .entries()
2695                    .get(&AttributeId::new(name).unwrap())
2696                    .unwrap()
2697                    .constraint,
2698                AttributeConstraint::TextChoices { values }
2699            );
2700        }
2701        for ordinal in [4, 5, 6, 9] {
2702            assert_eq!(
2703                linear.descriptor().inputs[ordinal].element_types(),
2704                &BTreeSet::from([ElementType::F32])
2705            );
2706        }
2707        assert_eq!(
2708            linear.descriptor().inputs[8].access(),
2709            TensorAccess::ReadWrite
2710        );
2711        assert_eq!(
2712            linear.descriptor().inputs[9].access(),
2713            TensorAccess::ReadWrite
2714        );
2715        assert_eq!(full.descriptor().inputs.len(), 9);
2716        assert_eq!(full.descriptor().version, ContractVersion::new(2, 0));
2717        assert_eq!(
2718            full.descriptor().provider.minimum_version,
2719            ContractVersion::new(2, 0)
2720        );
2721        assert_eq!(
2722            full.descriptor().resources.binding,
2723            ResourcePresenceRequirement::Required
2724        );
2725        assert_eq!(
2726            full.descriptor().inputs[8].access(),
2727            TensorAccess::ReadWrite
2728        );
2729    }
2730}