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