1use crate::{capture::*, ObservationPoint, ObservationSupportStatus, TensorAxis};
6use serde::{Deserialize, Serialize};
7use sha2::{Digest, Sha256};
8use std::collections::BTreeSet;
9
10pub const INTERVENTION_SCHEMA_VERSION: u32 = 1;
12pub const MAX_INTERVENTION_OPERATIONS: usize = 64;
14pub const MAX_INTERVENTION_PLAN_BYTES: u64 = 1024 * 1024;
16pub const MAX_INTERVENTION_PAYLOAD_BYTES: u64 = 512 * 1024;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
21#[serde(rename_all = "snake_case")]
22pub enum InterventionDtype {
23 Float32,
25 Float16,
27 Bfloat16,
29}
30
31#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
33#[serde(tag = "dtype", content = "values", rename_all = "snake_case")]
34pub enum InterventionValues {
35 Float32(Vec<f32>),
37 Float16(Vec<u16>),
39 Bfloat16(Vec<u16>),
41}
42
43impl InterventionValues {
44 pub fn dtype(&self) -> InterventionDtype {
46 match self {
47 Self::Float32(_) => InterventionDtype::Float32,
48 Self::Float16(_) => InterventionDtype::Float16,
49 Self::Bfloat16(_) => InterventionDtype::Bfloat16,
50 }
51 }
52 pub fn len(&self) -> usize {
54 match self {
55 Self::Float32(v) => v.len(),
56 Self::Float16(v) | Self::Bfloat16(v) => v.len(),
57 }
58 }
59 pub fn is_empty(&self) -> bool {
61 self.len() == 0
62 }
63 fn validate(&self) -> Result<(), CaptureError> {
64 let finite = match self {
65 Self::Float32(v) => v.iter().all(|v| v.is_finite()),
66 Self::Float16(v) => v.iter().all(|v| v & 0x7c00 != 0x7c00),
67 Self::Bfloat16(v) => v.iter().all(|v| v & 0x7f80 != 0x7f80),
68 };
69 require(finite, "intervention values must be finite")
70 }
71 fn bytes(&self) -> Result<u64, CaptureError> {
72 mul(
73 self.len() as u64,
74 if self.dtype() == InterventionDtype::Float32 {
75 4
76 } else {
77 2
78 },
79 )
80 }
81}
82
83#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
85pub struct InterventionTensor {
86 pub shape: Vec<u64>,
88 pub values: InterventionValues,
90}
91
92#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
94#[serde(rename_all = "snake_case")]
95pub enum RoutingScoreStage {
96 RawLogits,
98 TransformedScores,
100 RankingScores,
102}
103
104#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
106#[serde(rename_all = "snake_case")]
107pub enum RoutingScoring {
108 Softmax,
110 SelectedSoftmax,
112 Sigmoid,
114 SqrtSoftplus,
116}
117
118#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
120pub struct InterventionRoutingPolicy {
121 pub expert_count: u32,
123 pub top_k: u32,
125 pub scoring: RoutingScoring,
127 pub normalize_selected: bool,
129 pub normalization_epsilon: f32,
131 pub coefficient_scale: f32,
133 pub groups: u32,
135 pub selected_groups: u32,
137 pub learned_coefficient_scale: bool,
139 pub shared_experts: u32,
141}
142
143#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
145#[serde(rename_all = "snake_case")]
146pub enum InterventionStage {
147 Activation,
149 LogitsBeforeSampling,
151 RoutingBeforeDispatch,
153}
154
155#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
157#[serde(rename_all = "snake_case")]
158pub enum InterventionKind {
159 Zero,
161 Scale,
163 Mask,
165 Replace,
167 Add,
169 MaskLogits,
171 ExcludeExperts,
173 ZeroExpertContribution,
175 BiasRoutingScores,
177 ForceExperts,
179}
180
181#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
183pub struct InterventionPoint {
184 pub path: String,
186 pub node_id: String,
188 pub stage: InterventionStage,
190 pub axes: Vec<TensorAxis>,
192 pub dtypes: Vec<InterventionDtype>,
194 pub operations: Vec<InterventionKind>,
196 pub score_stages: Vec<RoutingScoreStage>,
198 pub prefill: ObservationSupportStatus,
200 pub decode: ObservationSupportStatus,
202 pub conditions: Vec<String>,
204 pub routing: Option<InterventionRoutingPolicy>,
206}
207
208impl InterventionPoint {
209 pub fn observation_geometry(&self) -> ObservationPoint {
211 ObservationPoint {
212 path: self.path.clone(),
213 node_id: self.node_id.clone(),
214 meaning: String::new(),
215 value_type: crate::ObservationValueType::Tensor,
216 dtype: crate::ObservationDtype::Floating,
217 axes: Some(self.axes.clone()),
218 prefill: true,
219 decode: true,
220 requirements: vec![],
221 position: crate::ObservationPosition::BeforeIntervention,
222 retained_bytes: None,
223 host_bytes: None,
224 }
225 }
226}
227
228#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
230pub struct InterventionDiscovery {
231 pub schema_version: u32,
233 pub artifact_identity: String,
235 #[serde(default)]
238 pub session_identity: Option<String>,
239 pub points: Vec<InterventionPoint>,
241}
242
243pub fn new_intervention_session_identity() -> String {
246 use std::sync::atomic::{AtomicU64, Ordering};
247 static NEXT: AtomicU64 = AtomicU64::new(0);
248 format!(
249 "intervention-session-{}-{}-{}",
250 std::process::id(),
251 std::time::SystemTime::now()
252 .duration_since(std::time::UNIX_EPOCH)
253 .map_or(0, |d| d.as_nanos()),
254 NEXT.fetch_add(1, Ordering::Relaxed)
255 )
256}
257
258#[derive(Debug, Clone, Default)]
260pub struct InterventionMechanisms {
261 pub operations: Vec<InterventionKind>,
263 pub dtypes: Vec<InterventionDtype>,
265 pub score_stages: Vec<RoutingScoreStage>,
267}
268
269#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
271#[serde(tag = "kind", rename_all = "snake_case")]
272pub enum InterventionAction {
273 Zero {
275 dtype: InterventionDtype,
277 },
278 Scale {
280 dtype: InterventionDtype,
282 factor: f32,
284 },
285 Mask {
287 dtype: InterventionDtype,
289 shape: Vec<u64>,
291 keep: Vec<bool>,
293 },
294 Replace {
296 tensor: InterventionTensor,
298 },
299 Add {
301 tensor: InterventionTensor,
303 },
304 MaskLogits {
306 dtype: InterventionDtype,
308 token_ids: Vec<u32>,
310 },
311 ExcludeExperts {
314 expert_ids: Vec<u32>,
316 },
317 ZeroExpertContribution {
321 expert_ids: Vec<u32>,
323 },
324 BiasRoutingScores {
326 stage: RoutingScoreStage,
328 expert_ids: Vec<u32>,
330 biases: Vec<f32>,
332 },
333 ForceExperts {
337 shape: [u64; 2],
339 expert_ids: Vec<u32>,
341 },
342}
343
344impl InterventionAction {
345 pub fn validate_activation_region(
347 &self,
348 dtype: InterventionDtype,
349 shape: &[u64],
350 ) -> Result<(), CaptureError> {
351 require(
352 self.dtype() == Some(dtype),
353 "intervention runtime dtype differs from exact declared dtype",
354 )?;
355 require(
356 !shape.is_empty() && shape.len() <= 32 && !shape.contains(&0),
357 "invalid activation region shape",
358 )?;
359 validate_activation_parameters(self, shape.last().and_then(|n| u32::try_from(*n).ok()))?;
360 validate_payload_shape(self, shape)
361 }
362 pub fn kind(&self) -> InterventionKind {
364 match self {
365 Self::Zero { .. } => InterventionKind::Zero,
366 Self::Scale { .. } => InterventionKind::Scale,
367 Self::Mask { .. } => InterventionKind::Mask,
368 Self::Replace { .. } => InterventionKind::Replace,
369 Self::Add { .. } => InterventionKind::Add,
370 Self::MaskLogits { .. } => InterventionKind::MaskLogits,
371 Self::ExcludeExperts { .. } => InterventionKind::ExcludeExperts,
372 Self::ZeroExpertContribution { .. } => InterventionKind::ZeroExpertContribution,
373 Self::BiasRoutingScores { .. } => InterventionKind::BiasRoutingScores,
374 Self::ForceExperts { .. } => InterventionKind::ForceExperts,
375 }
376 }
377 pub fn dtype(&self) -> Option<InterventionDtype> {
379 match self {
380 Self::Zero { dtype }
381 | Self::Scale { dtype, .. }
382 | Self::Mask { dtype, .. }
383 | Self::MaskLogits { dtype, .. } => Some(*dtype),
384 Self::Replace { tensor } | Self::Add { tensor } => Some(tensor.values.dtype()),
385 _ => None,
386 }
387 }
388 fn payload_bytes(&self) -> Result<u64, CaptureError> {
389 match self {
390 Self::Replace { tensor } | Self::Add { tensor } => {
391 add(mul(tensor.shape.len() as u64, 8)?, tensor.values.bytes()?)
392 }
393 Self::Mask { shape, keep, .. } => add(mul(shape.len() as u64, 8)?, keep.len() as u64),
394 Self::MaskLogits { token_ids, .. } => mul(token_ids.len() as u64, 4),
395 Self::ExcludeExperts { expert_ids }
396 | Self::ZeroExpertContribution { expert_ids }
397 | Self::ForceExperts { expert_ids, .. } => mul(expert_ids.len() as u64, 4),
398 Self::BiasRoutingScores {
399 expert_ids, biases, ..
400 } => mul(add(expert_ids.len() as u64, biases.len() as u64)?, 4),
401 _ => Ok(0),
402 }
403 }
404}
405
406#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
410#[serde(tag = "kind", rename_all = "snake_case")]
411pub enum InterventionEvidence {
412 None,
414 Preview {
416 max_elements: u64,
418 },
419 Summary,
421}
422
423#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
425pub struct InterventionOperation {
426 pub id: String,
428 pub target: String,
430 pub schedule: CaptureSchedule,
432 pub slices: Vec<CaptureSlice>,
434 pub action: InterventionAction,
436 pub evidence: InterventionEvidence,
438}
439
440#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
442pub struct InterventionPlan {
443 pub schema_version: u32,
445 pub operations: Vec<InterventionOperation>,
447}
448
449impl InterventionPlan {
450 pub fn none() -> Self {
452 Self {
453 schema_version: INTERVENTION_SCHEMA_VERSION,
454 operations: vec![],
455 }
456 }
457
458 pub fn admit(
461 self,
462 discovery: &InterventionDiscovery,
463 request: CaptureRequestShape,
464 session_id: &str,
465 ) -> Result<AdmittedInterventionPlan, CaptureError> {
466 require(
467 self.schema_version == INTERVENTION_SCHEMA_VERSION
468 && discovery.schema_version == INTERVENTION_SCHEMA_VERSION,
469 "unsupported intervention schema",
470 )?;
471 require(
472 !session_id.is_empty() && !discovery.artifact_identity.is_empty(),
473 "missing intervention session/source identity",
474 )?;
475 require(
476 discovery
477 .session_identity
478 .as_ref()
479 .is_some_and(|id| !id.is_empty()),
480 "intervention admission requires a realized backend session",
481 )?;
482 require(
483 request.batch > 0 && request.prompt_tokens > 0 && request.max_predictions > 0,
484 "empty intervention request geometry",
485 )?;
486 add(request.prompt_tokens, request.max_predictions)?;
487 mul(request.batch, request.prompt_tokens)?;
488 require(
489 self.operations.len() <= MAX_INTERVENTION_OPERATIONS,
490 "too many intervention operations",
491 )?;
492 let mut payload_bytes = 0;
493 let mut ids = BTreeSet::new();
494 let mut points = Vec::new();
495 for operation in &self.operations {
496 require(
497 !operation.id.is_empty() && operation.id.len() <= 128 && ids.insert(&operation.id),
498 "intervention IDs must be unique and contain 1..=128 bytes",
499 )?;
500 require(
501 operation.target.len() <= 1024,
502 "intervention target exceeds metadata bound",
503 )?;
504 payload_bytes = add(payload_bytes, operation.action.payload_bytes()?)?;
505 require(
506 payload_bytes <= MAX_INTERVENTION_PAYLOAD_BYTES,
507 "intervention payload exceeds hard bound",
508 )?;
509 let mut matching = discovery
510 .points
511 .iter()
512 .filter(|p| p.path == operation.target);
513 let point = matching
514 .next()
515 .ok_or_else(|| CaptureError::MissingPath(operation.target.clone()))?;
516 require(
517 matching.next().is_none(),
518 "ambiguous intervention target declaration",
519 )?;
520 validate_operation(operation, point, request)?;
521 points.push(point.clone());
522 }
523 for (index, (operation, point)) in self.operations.iter().zip(&points).enumerate() {
524 if point.routing.is_none() {
525 continue;
526 }
527 for previous in &self.operations[..index] {
528 if previous.target != operation.target {
529 continue;
530 }
531 for phase in [CapturePhase::Prefill, CapturePhase::Decode] {
532 let a = previous
533 .schedule
534 .count_and_last(phase, request.max_predictions)?;
535 let b = operation
536 .schedule
537 .count_and_last(phase, request.max_predictions)?;
538 if let (Some((ac, al)), Some((bc, bl))) = (a, b) {
539 let af = al - (ac - 1) * previous.schedule.every;
540 let bf = bl - (bc - 1) * operation.schedule.every;
541 require(
542 al < bf || bl < af,
543 "routing operations have potentially overlapping schedules",
544 )?;
545 }
546 }
547 }
548 }
549 let mut encoded = PlanCounter(0);
550 serde_json::to_writer(&mut encoded, &self)
551 .map_err(|_| CaptureError::Invalid("intervention plan exceeds encoded bound".into()))?;
552 let digest = Sha256::digest(
553 serde_json::to_vec(&(
554 &self,
555 &points,
556 request,
557 &discovery.artifact_identity,
558 &discovery.session_identity,
559 session_id,
560 ))
561 .map_err(|e| CaptureError::Invalid(e.to_string()))?,
562 );
563 let identity = format!(
564 "intervention-v1-{}",
565 digest
566 .iter()
567 .map(|byte| format!("{byte:02x}"))
568 .collect::<String>()
569 );
570 Ok(AdmittedInterventionPlan {
571 plan: self,
572 points,
573 request,
574 identity,
575 artifact_identity: discovery.artifact_identity.clone(),
576 session_id: session_id.into(),
577 })
578 }
579}
580
581struct PlanCounter(u64);
582impl std::io::Write for PlanCounter {
583 fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
584 self.0 = self
585 .0
586 .checked_add(bytes.len() as u64)
587 .filter(|n| *n <= MAX_INTERVENTION_PLAN_BYTES)
588 .ok_or_else(|| std::io::Error::other("intervention plan bound"))?;
589 Ok(bytes.len())
590 }
591 fn flush(&mut self) -> std::io::Result<()> {
592 Ok(())
593 }
594}
595
596#[derive(Debug, Clone)]
598pub struct AdmittedInterventionPlan {
599 plan: InterventionPlan,
600 points: Vec<InterventionPoint>,
601 request: CaptureRequestShape,
602 identity: String,
603 artifact_identity: String,
604 session_id: String,
605}
606
607impl AdmittedInterventionPlan {
608 pub fn identity(&self) -> &str {
610 &self.identity
611 }
612 pub fn artifact_identity(&self) -> &str {
614 &self.artifact_identity
615 }
616 pub fn session_id(&self) -> &str {
618 &self.session_id
619 }
620 pub fn plan(&self) -> &InterventionPlan {
622 &self.plan
623 }
624 pub fn points(&self) -> &[InterventionPoint] {
626 &self.points
627 }
628 pub fn request(&self) -> CaptureRequestShape {
630 self.request
631 }
632 pub fn is_empty(&self) -> bool {
634 self.plan.operations.is_empty()
635 }
636 pub fn validate_actual(
639 &self,
640 index: usize,
641 phase: CapturePhase,
642 prediction: u64,
643 shape: &[u64],
644 dtype: Option<InterventionDtype>,
645 ) -> Result<ResolvedCaptureSlice, CaptureError> {
646 let operation = self
647 .plan
648 .operations
649 .get(index)
650 .ok_or_else(|| CaptureError::Invalid("invalid admitted operation index".into()))?;
651 require(
652 prediction < self.request.max_predictions
653 && operation.schedule.includes(phase, prediction),
654 "inactive intervention operation",
655 )?;
656 require(
657 operation.action.dtype() == dtype,
658 "intervention runtime dtype differs from exact declared dtype",
659 )?;
660 let point = &self.points[index];
661 let geometry = point.observation_geometry();
662 self.request
663 .validate_actual(&geometry, phase, prediction, shape)?;
664 let slice = operation.resolve_slice(point, shape)?;
665 validate_payload_shape(&operation.action, &slice.shape)?;
666 Ok(slice)
667 }
668}
669
670impl InterventionOperation {
671 pub fn resolve_slice(
673 &self,
674 point: &InterventionPoint,
675 shape: &[u64],
676 ) -> Result<ResolvedCaptureSlice, CaptureError> {
677 resolve_slice(
678 &point.observation_geometry(),
679 &CaptureSelection {
680 id: self.id.clone(),
681 path: self.target.clone(),
682 schedule: self.schedule.clone(),
683 slices: self.slices.clone(),
684 transform: CaptureTransform::FullTensor,
685 },
686 shape,
687 )
688 }
689}
690
691fn validate_operation(
692 operation: &InterventionOperation,
693 point: &InterventionPoint,
694 request: CaptureRequestShape,
695) -> Result<(), CaptureError> {
696 require(
697 point.node_id.len() <= 1024 && !point.node_id.is_empty(),
698 "invalid intervention node identity",
699 )?;
700 require(
701 !point.axes.is_empty() && point.axes.len() <= 32,
702 "intervention requires declared bounded tensor rank",
703 )?;
704 require(
705 point.operations.contains(&operation.action.kind()),
706 "operation unsupported at intervention point",
707 )?;
708 require(
709 operation.schedule.every != 0
710 && operation
711 .schedule
712 .end_prediction
713 .is_none_or(|end| end > operation.schedule.first_prediction),
714 "invalid intervention schedule",
715 )?;
716 let mut axes = BTreeSet::new();
717 for axis in &point.axes {
718 require(
719 axes.insert(&axis.name),
720 "ambiguous intervention axis declaration",
721 )?;
722 }
723 let mut selected = BTreeSet::new();
724 for slice in &operation.slices {
725 require(
726 selected.insert(&slice.axis) && axes.contains(&slice.axis),
727 "duplicate or unknown intervention axis",
728 )?;
729 require(
730 slice.stride > 0 && slice.start < slice.end,
731 "intervention slices must be nonempty with positive stride",
732 )?;
733 if point.routing.is_some() {
734 require(
735 slice.axis == "token",
736 "routing permits token-row selection only",
737 )?;
738 }
739 if matches!(operation.action, InterventionAction::MaskLogits { .. }) {
740 require(
741 slice.axis != "vocabulary",
742 "logit masks require the complete vocabulary axis",
743 )?;
744 }
745 }
746 if let Some(dtype) = operation.action.dtype() {
747 require(
748 point.routing.is_none() && point.dtypes.contains(&dtype),
749 "activation dtype unsupported at target",
750 )?;
751 } else {
752 require(
753 point.routing.is_some(),
754 "routing action requires a pre-dispatch routing point",
755 )?;
756 }
757 match operation.evidence {
758 InterventionEvidence::Preview { max_elements } => require(
759 max_elements > 0 && max_elements <= 4096,
760 "intervention preview requires 1..=4096 elements",
761 )?,
762 InterventionEvidence::Summary => require(
763 point.routing.is_none(),
764 "routing evidence requires an ID/coefficient preview",
765 )?,
766 InterventionEvidence::None => (),
767 }
768 validate_action(&operation.action, point)?;
769 for (phase, status) in [
770 (CapturePhase::Prefill, &point.prefill),
771 (CapturePhase::Decode, &point.decode),
772 ] {
773 if let Some((_, last)) = operation
774 .schedule
775 .count_and_last(phase, request.max_predictions)?
776 {
777 if *status != ObservationSupportStatus::Supported {
778 return Err(CaptureError::Unsupported(format!(
779 "intervention {} {phase:?}: {status:?}",
780 point.path
781 )));
782 }
783 let geometry = point.observation_geometry();
784 if let Some(shape) = request.resolve(&geometry, phase, last)? {
785 let slice = operation.resolve_slice(point, &shape)?;
786 validate_payload_shape(&operation.action, &slice.shape)?;
787 }
788 }
789 }
790 Ok(())
791}
792
793fn validate_activation_parameters(
794 action: &InterventionAction,
795 vocabulary: Option<u32>,
796) -> Result<(), CaptureError> {
797 use InterventionAction as A;
798 match action {
799 A::Scale { dtype, factor } => {
800 let maximum = match dtype {
801 InterventionDtype::Float16 => 65504.0,
802 InterventionDtype::Bfloat16 => f32::from_bits(0x7f7f0000),
803 InterventionDtype::Float32 => f32::MAX,
804 };
805 require(
806 factor.is_finite() && factor.abs() <= maximum,
807 "scale must be finite and representable in the target dtype",
808 )?;
809 }
810 A::Replace { tensor } | A::Add { tensor } => {
811 require(
812 !tensor.shape.is_empty() && tensor.shape.len() <= 32 && !tensor.shape.contains(&0),
813 "invalid replacement tensor rank/extents",
814 )?;
815 require(
816 elements(&tensor.shape)? == tensor.values.len() as u64,
817 "replacement tensor is incomplete or has malformed shape",
818 )?;
819 tensor.values.validate()?;
820 }
821 A::Mask { shape, keep, .. } => {
822 require(
823 !shape.is_empty() && shape.len() <= 32 && !shape.contains(&0),
824 "invalid mask rank/extents",
825 )?;
826 require(
827 elements(shape)? == keep.len() as u64,
828 "mask shape differs from complete keep values",
829 )?;
830 }
831 A::MaskLogits { token_ids, .. } => {
832 let vocabulary = vocabulary.ok_or_else(|| {
833 CaptureError::Unsupported("logit mask requires known vocabulary extent".into())
834 })?;
835 validate_ids(token_ids, vocabulary)?;
836 require(
837 token_ids.len() < vocabulary as usize,
838 "logit mask cannot exclude the whole vocabulary",
839 )?;
840 }
841 A::Zero { .. } => (),
842 _ => {
843 return Err(CaptureError::Invalid(
844 "routing action cannot replace an activation".into(),
845 ))
846 }
847 }
848 Ok(())
849}
850
851fn validate_action(
852 action: &InterventionAction,
853 point: &InterventionPoint,
854) -> Result<(), CaptureError> {
855 use InterventionAction as A;
856 if action.dtype().is_some() {
857 let vocabulary = if matches!(action, A::MaskLogits { .. }) {
858 require(
859 point.stage == InterventionStage::LogitsBeforeSampling,
860 "logit masking requires final logits",
861 )?;
862 point
863 .axes
864 .iter()
865 .find(|axis| axis.name == "vocabulary")
866 .and_then(|axis| match axis.dimension {
867 crate::SymbolicDimension::Known(n) => u32::try_from(n).ok(),
868 _ => None,
869 })
870 } else {
871 None
872 };
873 return validate_activation_parameters(action, vocabulary);
874 }
875
876 match action {
877 A::ExcludeExperts { expert_ids }
878 | A::ZeroExpertContribution { expert_ids }
879 | A::BiasRoutingScores { expert_ids, .. }
880 | A::ForceExperts { expert_ids, .. } => {
881 let policy = point
882 .routing
883 .as_ref()
884 .ok_or_else(|| CaptureError::Invalid("missing routing policy".into()))?;
885 require(
886 point.stage == InterventionStage::RoutingBeforeDispatch,
887 "routing control requires pre-dispatch target",
888 )?;
889 require(
890 policy.expert_count > 0
891 && policy.top_k > 0
892 && policy.top_k <= policy.expert_count
893 && policy.groups > 0
894 && policy.expert_count % policy.groups == 0
895 && policy.selected_groups > 0
896 && policy.selected_groups <= policy.groups
897 && policy.top_k
898 <= policy.selected_groups * (policy.expert_count / policy.groups)
899 && policy.normalization_epsilon.is_finite()
900 && policy.normalization_epsilon >= 0.0
901 && policy.coefficient_scale.is_finite()
902 && policy.coefficient_scale > 0.0,
903 "invalid routing declaration",
904 )?;
905 if let A::ForceExperts { shape, .. } = action {
906 require(
907 shape[0] > 0
908 && shape[1] == policy.top_k as u64
909 && elements(shape)? == expert_ids.len() as u64,
910 "forced route shape must be [selected_token_rows, top_k]",
911 )?;
912 for row in expert_ids.chunks(policy.top_k as usize) {
913 validate_ids(row, policy.expert_count)?;
914 let width = policy.expert_count / policy.groups;
915 let groups: BTreeSet<_> = row.iter().map(|id| id / width).collect();
916 require(
917 groups.len() <= policy.selected_groups as usize,
918 "forced IDs exceed architecture's selected-group count",
919 )?;
920 }
921 } else {
922 validate_ids(expert_ids, policy.expert_count)?;
923 }
924 if let A::ExcludeExperts { .. } = action {
925 require(
926 policy.expert_count as usize - expert_ids.len() >= policy.top_k as usize,
927 "excluded experts make top-k infeasible",
928 )?;
929 let width = policy.expert_count / policy.groups;
933 let mut available = vec![width; policy.groups as usize];
934 for id in expert_ids {
935 available[(id / width) as usize] -= 1;
936 }
937 available.sort_unstable();
938 require(
939 available
940 .iter()
941 .take(policy.selected_groups as usize)
942 .sum::<u32>()
943 >= policy.top_k,
944 "exclusion cannot guarantee top-k within selected groups",
945 )?;
946 }
947 if let A::BiasRoutingScores { stage, biases, .. } = action {
948 require(
949 point.score_stages.contains(stage),
950 "router bias stage unsupported",
951 )?;
952 require(
953 biases.len() == expert_ids.len() && biases.iter().all(|b| b.is_finite()),
954 "router bias must have one finite value per expert",
955 )?;
956 }
957 }
958 _ => unreachable!("activation actions validated above"),
959 }
960 Ok(())
961}
962
963fn validate_payload_shape(
964 action: &InterventionAction,
965 selected: &[u64],
966) -> Result<(), CaptureError> {
967 let expected: Option<&[u64]> = match action {
968 InterventionAction::Mask { shape, .. } => Some(shape),
969 InterventionAction::Replace { tensor } | InterventionAction::Add { tensor } => {
970 Some(&tensor.shape)
971 }
972 InterventionAction::ForceExperts { shape, .. } => Some(shape),
973 _ => None,
974 };
975 require(!selected.contains(&0), "empty intervention selection")?;
976 require(
977 expected.is_none_or(|expected| expected == selected),
978 "intervention payload must exactly match selected shape; broadcasting is prohibited",
979 )
980}
981
982fn validate_ids(ids: &[u32], count: u32) -> Result<(), CaptureError> {
983 require(!ids.is_empty(), "intervention ID list must be nonempty")?;
984 let mut unique = BTreeSet::new();
985 require(
986 ids.iter().all(|id| *id < count && unique.insert(*id)),
987 "intervention IDs are duplicate or out of range",
988 )
989}
990
991fn require(condition: bool, message: &str) -> Result<(), CaptureError> {
992 if condition {
993 Ok(())
994 } else {
995 Err(CaptureError::Invalid(message.into()))
996 }
997}
998
999#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
1001#[serde(tag = "kind", rename_all = "snake_case")]
1002pub enum InterventionOutcome {
1003 Inactive,
1005 Applied,
1007 Missing,
1009 Failed {
1011 message: String,
1013 },
1014}
1015
1016#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1018pub struct InterventionRecord {
1019 pub schema_version: u32,
1021 pub plan_id: String,
1023 pub operation_id: String,
1025 pub target: String,
1027 pub node_id: String,
1029 pub phase: CapturePhase,
1031 pub prediction_index: u64,
1033 pub outcome: InterventionOutcome,
1035 pub evidence: Vec<CaptureRecord>,
1038 pub charged: CaptureUsage,
1040}
1041
1042pub trait InterventionBackend: CaptureBackend {
1045 fn intervention_dtype(&self, tensor: &Self::Tensor) -> Result<InterventionDtype, Self::Error>;
1047 fn validate_intervention_geometry(
1049 &self,
1050 source: &[u64],
1051 slice: &ResolvedCaptureSlice,
1052 ) -> Result<(), CaptureError>;
1053 fn select_region(
1055 &mut self,
1056 tensor: &Self::Tensor,
1057 slice: &ResolvedCaptureSlice,
1058 ) -> Result<Self::Tensor, Self::Error>;
1059 fn update_region(
1061 &mut self,
1062 tensor: &Self::Tensor,
1063 slice: &ResolvedCaptureSlice,
1064 replacement: &Self::Tensor,
1065 ) -> Result<Self::Tensor, Self::Error>;
1066 fn zeros(
1068 &mut self,
1069 shape: &[u64],
1070 dtype: InterventionDtype,
1071 ) -> Result<Self::Tensor, Self::Error>;
1072 fn scale(&mut self, value: &Self::Tensor, factor: f32) -> Result<Self::Tensor, Self::Error>;
1074 fn fill_masked(
1076 &mut self,
1077 value: &Self::Tensor,
1078 keep: &[bool],
1079 fill: f32,
1080 ) -> Result<Self::Tensor, Self::Error>;
1081 fn realize_tensor(&mut self, tensor: &InterventionTensor) -> Result<Self::Tensor, Self::Error>;
1083 fn add(
1085 &mut self,
1086 left: &Self::Tensor,
1087 right: &Self::Tensor,
1088 ) -> Result<Self::Tensor, Self::Error>;
1089 fn fill_columns(
1091 &mut self,
1092 value: &Self::Tensor,
1093 ids: &[u32],
1094 fill: f32,
1095 ) -> Result<Self::Tensor, Self::Error>;
1096}
1097
1098pub trait InterventionEstimator: Send + Sync {
1101 fn validate_geometry(
1103 &self,
1104 source: &[u64],
1105 slice: &ResolvedCaptureSlice,
1106 ) -> Result<(), CaptureError>;
1107 fn capture_usage(
1109 &self,
1110 source: &[u64],
1111 selection: &CaptureSelection,
1112 slice: &ResolvedCaptureSlice,
1113 ) -> Result<CaptureUsage, CaptureError>;
1114 fn original_route_usage(
1117 &self,
1118 policy: &InterventionRoutingPolicy,
1119 rows: u64,
1120 ) -> Result<CaptureUsage, CaptureError>;
1121}
1122
1123#[cfg(test)]
1124mod tests;