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