Skip to main content

katgpt_types/
enums.rs

1//! Small configuration enums and feature-config structs.
2
3// Shared configuration, RNG, and math utilities.
4// Superset of types from both katgpt-rs and riir-engine projects.
5
6// ---------------------------------------------------------------------------
7// Enums
8// ---------------------------------------------------------------------------
9
10/// Adaptive depth tier mapping to layer count (Plan 284).
11/// Reuses ThermalPath naming convention from FlashAR Consensus (Plan 166).
12///
13/// | Tier     | Layers     | When                              |
14/// |----------|------------|-----------------------------------|
15/// | Plasma   | 1          | High entropy, easy positions      |
16/// | Hot      | 2          | Medium entropy, standard tactics  |
17/// | Warm     | all        | Low entropy, complex positions    |
18/// | Cold     | all+verify | Critical, full verification       |
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20#[repr(u8)]
21pub enum DepthTier {
22    /// Easy positions: empty board, forced moves. 1 layer.
23    Plasma = 0,
24    /// Moderate: standard tactics. 2 layers.
25    Hot = 1,
26    /// Complex: all layers + spot-check verification.
27    Warm = 2,
28    /// Critical: all layers + full verification.
29    Cold = 3,
30}
31
32impl DepthTier {
33    /// Returns the maximum number of transformer layers to execute for this tier.
34    pub fn max_layers(&self, total_layers: usize) -> usize {
35        match self {
36            DepthTier::Plasma => 1.min(total_layers),
37            DepthTier::Hot => 2.min(total_layers),
38            DepthTier::Warm => total_layers,
39            DepthTier::Cold => total_layers,
40        }
41    }
42}
43
44/// Attention mode for HLA (Higher-order Linear Attention).
45///
46/// - `Standard`: SDPA with KV cache (default, backward-compatible).
47/// - `Hla`: Symmetric second-order linear attention — O(1) per-token memory.
48/// - `Ahla`: Asymmetric second-order linear attention — lower state cost than symmetric.
49#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
50#[repr(u8)]
51pub enum HlaMode {
52    #[default]
53    Standard,
54    /// Symmetric second-order: SK, CQV, mQ accumulators.
55    Hla,
56    /// Asymmetric second-order: PKV, mK accumulators.
57    Ahla,
58}
59
60/// Attention mode for forward passes.
61///
62/// - `Causal`: Standard autoregressive — only attend to positions ≤ current (default).
63/// - `Bidirectional`: Attend to ALL positions — used for dLLM masked prediction (Plan 066).
64/// - `BlockCausal`: Bidirectional within current block, causal across blocks — D2F student.
65/// - `SpKv`: SP-KV self-pruned key-value attention (Plan 070).
66/// - `SpKvQuant`: SP-KV + Quantized KV fusion (Plan 070 Phase 3, Task T12).
67#[repr(u8)]
68#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
69pub enum AttentionMode {
70    #[default]
71    Causal,
72    /// Full bidirectional: all positions see all positions (teacher mode).
73    Bidirectional,
74    /// Block-causal: bidirectional within block, causal across blocks (student mode).
75    BlockCausal,
76    /// SP-KV self-pruned key-value attention (Plan 070).
77    /// Learns which KV pairs to retain via utility prediction.
78    /// Gate bias = log(u) during training, 0|-inf during inference.
79    SpKv,
80    /// SP-KV + Quantized KV fusion (Plan 070 Phase 3, Task T12).
81    /// Selective write (SP-KV utility gating) + lossy quantize (any QuantizedKVCache backend).
82    /// Two-stage compression: only useful KV pairs kept, those compressed to 2-4 bits/coord.
83    SpKvQuant,
84    /// DashAttention: adaptive sparse hierarchical attention via α-entmax routing (Plan 106).
85    /// Replaces fixed-budget top-k block selection with adaptive support selection.
86    /// Learned chunk summaries via head_cls vectors.
87    DashAttn,
88}
89
90/// Model architecture selector for forward pass dispatch.
91#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
92#[repr(u8)]
93pub enum ModelArchitecture {
94    #[default]
95    Generic,
96    Gemma2,
97    Llama,
98    /// Hybrid DeltaNet/Attention model (e.g., Qwen 3.5, Kimi Linear).
99    /// Uses per-layer config to determine DeltaNet vs standard attention.
100    /// Plan 182: Luce Megakernel Distill — DeltaNet GPU Inference.
101    #[cfg(feature = "deltanet_inference")]
102    QwenDeltaNet,
103}
104
105/// Attention projection configuration.
106/// Controls whether K and V projections share weights (Q-K=V tying).
107#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
108#[repr(u8)]
109pub enum AttentionProjection {
110    /// Standard Q, K, V (3 projections, full KV cache)
111    #[default]
112    Full,
113    /// Q-K=V: K and V share projection (2 projections, K-only cache).
114    /// 50% KV cache reduction, ~3% perplexity cost.
115    /// Post-hoc weight merging: W_kv = (W_k + W_v) / 2.
116    SharedKV,
117}
118
119/// KV cache layout (derived from AttentionProjection).
120#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121#[repr(u8)]
122pub enum CacheLayout {
123    /// Store both K and V (standard)
124    KV,
125    /// Store K only, V = K at read (SharedKV)
126    K,
127}
128
129/// Weight storage dtype (affects loading and dequantization).
130#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
131#[repr(u8)]
132pub enum WeightDtype {
133    #[default]
134    F32,
135    F16,
136    BF16,
137}
138
139// ---------------------------------------------------------------------------
140// Delta Routing (Plan 097, Research 061)
141// ---------------------------------------------------------------------------
142
143/// Delta routing mode — cross-layer information flow via delta vectors.
144/// Research 061: Delta Attention Residuals (Plan 097).
145///
146/// Kept compiled even when `delta_routing` is off so config round-trips
147/// serialize identically across feature sets. Reachable via `Config` defaults
148/// once the routing backend lands.
149#[allow(dead_code)]
150#[repr(u8)]
151#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
152pub enum DeltaRoutingMode {
153    /// No delta routing (default).
154    #[default]
155    Off,
156    /// Delta Block: accumulate deltas within blocks of `block_size` layers.
157    /// B+1 sources per routing decision. ~20% throughput overhead.
158    DeltaBlock,
159    /// Delta Attention Residuals: per-sublayer delta routing.
160    /// 2L sources. 69% throughput reduction at L=36. Use only for research.
161    DeltaAttnRes,
162}
163
164/// Configuration for delta routing (Plan 097, Research 061).
165///
166/// Fields ordered by descending alignment to minimize padding:
167/// usize (8B) → repr(u8) enum (1B) — 16 bytes total, no wasted padding.
168#[allow(dead_code)]
169#[derive(Clone, Copy, Debug)]
170pub struct DeltaRoutingConfig {
171    /// Block size for DeltaBlock mode (number of layers per block).
172    /// Default: 4. Paper recommends B=4.
173    pub block_size: usize,
174    /// Routing mode.
175    pub mode: DeltaRoutingMode,
176}
177
178impl Default for DeltaRoutingConfig {
179    fn default() -> Self {
180        Self {
181            block_size: 4,
182            mode: DeltaRoutingMode::Off,
183        }
184    }
185}
186
187// ---------------------------------------------------------------------------
188// DeltaNet Inference (Plan 182: Luce Megakernel Distill)
189// ---------------------------------------------------------------------------
190
191/// Per-layer type for hybrid DeltaNet/Attention models.
192/// Each layer is either a standard attention layer or a DeltaNet recurrent layer.
193#[cfg(feature = "deltanet_inference")]
194#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
195#[repr(u8)]
196pub enum DeltaNetLayerType {
197    /// Standard multi-head attention with KV cache.
198    #[default]
199    Attention,
200    /// DeltaNet linear recurrent layer (fast recurrent update, no KV cache needed).
201    DeltaNet,
202}
203
204// DeltaRoutingConfig::delta_block / is_enabled are intended for the
205// delta_routing backend (Plan 097) which is still being wired up. Silence
206// dead-code until callers land.
207#[allow(dead_code)]
208impl DeltaRoutingConfig {
209    pub fn delta_block(block_size: usize) -> Self {
210        Self {
211            mode: DeltaRoutingMode::DeltaBlock,
212            block_size,
213        }
214    }
215
216    pub fn is_enabled(&self) -> bool {
217        self.mode != DeltaRoutingMode::Off
218    }
219}
220
221// ---------------------------------------------------------------------------
222// DashAttention Config (Plan 106, Research 68)
223// ---------------------------------------------------------------------------
224
225/// Configuration for DashAttention adaptive sparse hierarchical attention.
226/// Controls α-entmax routing, chunk summarization, and routing bias.
227///
228/// Fields ordered by descending alignment to minimize padding:
229/// usize (8B) → f32 (4B) → bool (1B) — 24 bytes total, no wasted padding.
230#[derive(Clone, Copy, Debug)]
231pub struct DashAttnConfig {
232    /// Chunk size for block-level attention (default: 64).
233    pub chunk_size: usize,
234    /// α parameter for entmax. Only α=1.5 supported (quadratic, closed-form).
235    pub alpha: f32,
236    /// Scaling factor γ applied to chunk logits before entmax (default: 1.0).
237    pub scaling_factor: f32,
238    /// Prior strength σ for routing bias (default: 1e6, weak prior).
239    pub sigma: f32,
240    /// Whether to estimate diagonal attention contribution (default: true).
241    /// Tail-packed after f32 group to avoid bool-between-f32 padding.
242    pub estimate_diagonal: bool,
243}
244
245impl Default for DashAttnConfig {
246    fn default() -> Self {
247        Self {
248            chunk_size: 64,
249            alpha: 1.5,
250            scaling_factor: 1.0,
251            sigma: 1e6,
252            estimate_diagonal: true,
253        }
254    }
255}
256
257// ---------------------------------------------------------------------------
258// RTPurbo Retrieval Head Sparse Decode (Plan 126, Research 86)
259// ---------------------------------------------------------------------------
260
261/// Head role classification for RTPurbo sparse decode.
262///
263/// Only ~15% of attention heads ("retrieval heads") need full long-context access.
264/// The remaining ~85% ("local heads") attend only to local context + attention sinks.
265#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
266#[repr(u8)]
267pub enum RetrievalHeadRole {
268    /// Local head — sliding window + sink tokens only, no full KV scan.
269    #[default]
270    Local,
271    /// Retrieval head — low-dim projection + dynamic top-p token selection.
272    Retrieval,
273}
274
275/// Configuration for RTPurbo retrieval head sparse decode.
276///
277/// Feature gate: `rt_turbo` (opt-in, requires `dash_attn`).
278/// Adds head-wise retrieval/local classification + dynamic top-p token selection
279/// for decode-phase sparse attention. Complements DashAttention's α-entmax block
280/// routing with per-head specialization.
281///
282/// Must pass 6/6 GOAT proofs before default-on promotion.
283///
284/// Fields ordered by descending alignment to minimize padding:
285/// usize (8B) → f32 (4B) → CalibrationMode (1B) — no padding between groups.
286///
287/// # Calibration mode (Plan 358)
288///
289/// [`CalibrationMode::AttentionMass`] is the default (cheaper: 1 forward pass).
290/// [`CalibrationMode::CausalNecessity`] is opt-in — strictly stronger on
291/// workloads with correlated bystanders but ~10–100× more expensive to
292/// calibrate. See `calibrate_from_causal_scores` in `rt_turbo::calibration`.
293#[derive(Clone, Copy, Debug)]
294pub struct RtTurboConfig {
295    /// Low-dimensional projection size for pre-RoPE scoring (default: 16).
296    /// Paper ablation: dim=16 is the sweet spot for low-frequency retrieval.
297    pub low_dim: usize,
298    /// Sliding window size for local heads (default: 8192).
299    pub sliding_window: usize,
300    /// Number of attention sink tokens always retained for local heads (default: 4).
301    pub sink_tokens: usize,
302    /// Block size for block-level top-p variant (default: 64).
303    /// Should match `DashAttnConfig::chunk_size` for consistent routing.
304    pub block_size: usize,
305    /// Fraction of heads classified as retrieval heads (default: 0.15).
306    /// Paper ablation: 15% is optimal balance of accuracy vs sparsity.
307    pub retrieval_head_ratio: f32,
308    /// Cumulative attention mass threshold for dynamic top-p selection (default: 0.9).
309    /// Paper ablation: top-p=0.9 preserves >93% attention mass at 97% sparsity.
310    pub top_p: f32,
311    /// Which score semantics to use for head calibration (Plan 358). Default:
312    /// `AttentionMass` (cheaper). `CausalNecessity` is opt-in — strictly
313    /// stronger on bystander-heavy workloads but ~10–100× more expensive.
314    /// `AdaptiveCausal` (Proposal 004) is opt-in — cheap-proxy escalate,
315    /// unvalidated, requires per-head OV norms from the caller.
316    pub calibration_mode: CalibrationMode,
317}
318
319/// Head-calibration score source (Plan 358, Research 362).
320///
321/// `#[repr(u8)]` for sync-friendly 1-byte representation.
322#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
323#[repr(u8)]
324pub enum CalibrationMode {
325    /// Observational needle attention-mass (RTPurbo Plan 126 default).
326    /// Cheaper: a single forward pass + per-head mass scan.
327    #[default]
328    AttentionMass = 0,
329    /// Causal necessity via activation/path patching IE score (Plan 358).
330    /// Strictly stronger — excludes correlated bystanders — but requires
331    /// `n_heads × n_calibration_samples` patched forward passes. Requires the
332    /// `causal_head_importance` feature on the consuming crate.
333    CausalNecessity = 1,
334    /// Adaptive cheap-proxy escalate (Proposal 004 — OUR INVENTION, not from
335    /// HydraHead). Uses an OV-circuit proxy (`attention_mass / ||OV_out||`)
336    /// to detect bystander suspects, then escalates to Plan 358's causal
337    /// patching only on those `k` suspects instead of all `n_heads`. Pays zero
338    /// patched forwards when there are no bystanders (degenerates to
339    /// `AttentionMass`). Requires the `adaptive_causal_calibration` feature.
340    ///
341    /// **UNVALIDATED.** Promotion to default is blocked on G1 (proxy precision)
342    /// and G2 (cost reduction), both deferred to riir-engine. Unlike the other
343    /// two modes, the caller must supply per-head OV output norms (from a real
344    /// transformer forward) — see `calibrate_from_adaptive_causal`.
345    AdaptiveCausal = 2,
346}
347
348impl Default for RtTurboConfig {
349    fn default() -> Self {
350        Self {
351            low_dim: 16,
352            sliding_window: 8192,
353            sink_tokens: 4,
354            block_size: 64,
355            retrieval_head_ratio: 0.15,
356            top_p: 0.9,
357            calibration_mode: CalibrationMode::default(),
358        }
359    }
360}
361
362// ---------------------------------------------------------------------------
363// LT2 Looped Inference (Plan 108, Research 73)
364// ---------------------------------------------------------------------------
365
366/// Looped transformer mode — weight-shared layer repetition.
367#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
368#[repr(u8)]
369pub enum LoopMode {
370    /// Standard single-pass (no looping).
371    #[default]
372    None,
373    /// Weight-shared looping: same layers applied T times.
374    /// Effective depth = n_layer × loop_count.
375    WeightShared { loop_count: usize },
376    /// Training-free loop: ODE-refined sub-stepping over a window of layers.
377    /// No extra parameters — pure inference-time retrofit (Plan 136).
378    TrainingFree,
379}
380
381/// Hybrid attention pattern for looped inference.
382/// Controls which layers use full SDPA vs linear attention.
383#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
384#[repr(u8)]
385pub enum HybridPattern {
386    /// All layers use the same attention mode.
387    #[default]
388    Uniform,
389    /// Depth-level interleave: every Nth layer uses full SDPA.
390    /// e.g., Interleave { full_ratio: 5 } = every 5th layer is full.
391    /// Paper optimal: 1:4 ratio (full_ratio=5).
392    Interleave { full_ratio: usize },
393    /// Bookend: first and last layers are full, middle is linear.
394    Bookend,
395}
396
397/// Loop stability mode for weight-shared looped inference (Plan 428).
398///
399/// Parameter-free architectural fixes for T-pass loop stability, validated
400/// via the §3.6 defend-wrong PoC benchmark.
401///
402/// **Only `InterLoopNorm` ships** — the PoC proved it's the sole fix that
403/// controls residual norm growth. FLA-res (direct residual addition of
404/// `prev_h` at every layer) caused catastrophic norm explosion (~2.2B× at
405/// T=12), and Attention Injection was a no-op for single-position attention
406/// (softmax of 1 element = 1.0, so Q doesn't affect the output). Both were
407/// dropped per the defend-wrong verdict.
408#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
409#[repr(u8)]
410pub enum LoopStabilityMode {
411    /// No inter-loop stabilization (byte-identical to pre-Plan-428 behavior).
412    #[default]
413    None,
414    /// Inter-loop RMSNorm: normalize the hidden state between loop iterations
415    /// (tau > 0), before the inner layer pass. PoC: norm ratio 3.34× vs
416    /// baseline 11.19×, KL 0.0008, step-size trend converging (14.9 → 2.05).
417    InterLoopNorm,
418}
419
420/// Head-specific sigmoid gate after SDPA, before Wo.
421/// Zero-initialized → starts at sigmoid(0) = 0.5 (neutral multiplicative identity).
422#[derive(Clone)]
423pub struct SdpaOutputGate {
424    /// Gate weights: [n_heads * head_dim, dim].
425    /// Zero-init so gate starts at sigmoid(0) = 0.5.
426    pub w_gate: Vec<f32>,
427}
428
429impl SdpaOutputGate {
430    /// Allocate zeroed gate weights.
431    pub fn new(n_heads: usize, head_dim: usize, dim: usize) -> Self {
432        Self {
433            w_gate: vec![0.0; n_heads * head_dim * dim],
434        }
435    }
436
437    /// Apply sigmoid-gated projection to attention output.
438    ///
439    /// Computes: `gate[i] = sigmoid(W_gate[i] · attn_out)`, then `attn_out[i] *= gate[i]`.
440    /// Zero-init weights produce sigmoid(0) = 0.5 for all (neutral half-pass).
441    /// Paper reference: +0.3–0.5 avg points on zero-shot benchmarks.
442    pub fn forward(&self, attn_out: &mut [f32], dim: usize, temp: &mut [f32]) {
443        let n = attn_out.len();
444        debug_assert!(temp.len() >= n, "temp buffer too small");
445        debug_assert!(self.w_gate.len() >= n * dim, "gate weights too small");
446
447        // Step 1: Compute gate signal = sigmoid(W_gate @ attn_out)
448        // Batch matvec then batch sigmoid avoids per-element loop overhead
449        crate::simd::simd_matvec(temp, &self.w_gate, attn_out, n, dim);
450
451        // SIMD sigmoid: temp = -temp, exp, then 1/(1+exp)
452        crate::simd::simd_scale_inplace(&mut temp[..n], -1.0);
453        crate::simd::simd_exp_inplace(&mut temp[..n]);
454        crate::simd::simd_add_scalar_inplace(&mut temp[..n], 1.0);
455        // temp now = 1 + exp(-x), invert: temp = 1/temp = sigmoid
456        crate::simd::simd_reciprocal_inplace(&mut temp[..n]);
457
458        // Step 2: Apply gate elementwise via SIMD scale-mul (fused)
459        // attn_out[i] *= temp[i] is element-wise multiply
460        // Use simd_scale_mul_inplace with scale=1.0: attn[i] = temp[i] * attn[i] * 1.0
461        crate::simd::simd_scale_mul_inplace(attn_out, &temp[..n], 1.0);
462    }
463}
464
465/// Per-loop residual scaling gate.
466/// h^(τ) = h̃^(τ) + ρ_τ ⊙ h^(τ-1)
467/// Zero-init so first iteration is h̃^(1) (no residual from "previous").
468#[derive(Clone)]
469pub struct ResidualGate {
470    /// Per-loop gates: [loop_count, dim].
471    /// Each ρ_τ is element-wise, zero-init.
472    pub gates: Vec<f32>,
473}
474
475impl ResidualGate {
476    /// Allocate zeroed residual gates.
477    pub fn new(loop_count: usize, dim: usize) -> Self {
478        Self {
479            gates: vec![0.0; loop_count * dim],
480        }
481    }
482
483    /// Deterministic loop-stable residual gates (Plan 483 T2.1, §3.5 path 2).
484    ///
485    /// Sets gates to a constant `decay` factor for τ > 0, enabling information
486    /// carry-forward between T passes. The first loop (τ=0) has zero gates
487    /// (no previous state to carry forward).
488    ///
489    /// This is a **modelless** construction — no training, no gradient descent.
490    /// The `decay` factor controls the trade-off between carry-forward strength
491    /// and stability. Conservative values (0.1–0.3) are safe for most weight
492    /// matrices; larger values risk divergence.
493    ///
494    /// Rationale: the zero-init default (`new()`) makes every T-pass effectively
495    /// independent — no hidden state carries forward between loops. This
496    /// undermines the LT2 paper's "effective depth T×n_layer" claim. A non-zero
497    /// deterministic gate restores the residual connection across loops without
498    /// requiring trained gate parameters.
499    ///
500    /// # Arguments
501    /// * `loop_count` - Number of T-passes (T)
502    /// * `dim` - Hidden dimension (n_embd)
503    /// * `decay` - Constant gate value for τ > 0 (e.g., 0.1 = conservative)
504    #[inline]
505    pub fn new_loop_stable(loop_count: usize, dim: usize, decay: f32) -> Self {
506        let mut gates = vec![0.0f32; loop_count * dim];
507        // τ=0: zero (no previous state). τ>0: constant decay.
508        for tau in 1..loop_count {
509            let offset = tau * dim;
510            gates[offset..offset + dim].fill(decay);
511        }
512        Self { gates }
513    }
514
515    /// Deterministic loop-stable residual gates with exponential decay
516    /// (Plan 483 T2.1, §3.5 path 2 variant).
517    ///
518    /// ρ_τ = `base`^(τ-1) for τ > 0 — later loops contribute exponentially less.
519    /// This schedule mirrors the spectral-radius-based stabilization where the
520    /// residual contribution decays as the hidden state converges.
521    ///
522    /// # Arguments
523    /// * `loop_count` - Number of T-passes (T)
524    /// * `dim` - Hidden dimension (n_embd)
525    /// * `base` - Decay base (e.g., 0.5 → ρ_1=1.0, ρ_2=0.5, ρ_3=0.25, ...)
526    #[inline]
527    pub fn new_loop_stable_exp_decay(loop_count: usize, dim: usize, base: f32) -> Self {
528        let mut gates = vec![0.0f32; loop_count * dim];
529        for tau in 1..loop_count {
530            let offset = tau * dim;
531            let val = base.powi((tau - 1) as i32);
532            gates[offset..offset + dim].fill(val);
533        }
534        Self { gates }
535    }
536}
537
538// ---------------------------------------------------------------------------
539// SR²AM Configurator Bandit (Plan 112, Research 076)
540// ---------------------------------------------------------------------------
541
542/// SR²AM Configurator decision — learned per-turn planning regulation.
543///
544/// The configurator selects one of these arms per inference turn based on
545/// context (domain + entropy bin). UCB1 balances exploration vs exploitation.
546///
547/// - `PlanNew`: reset tree, full budget allocation (high uncertainty, new sub-problem)
548/// - `PlanExtend`: keep tree, extend depth by one level (moderate uncertainty, continuing)
549/// - `PlanSkip`: skip tree search, direct token sampling (low uncertainty, confident)
550#[cfg(feature = "sr2am_configurator")]
551#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
552#[repr(u8)]
553pub enum PlanningDecision {
554    /// Reset tree, full budget allocation (high uncertainty, new sub-problem).
555    PlanNew,
556    /// Keep tree, extend depth by one level (moderate uncertainty, continuing).
557    PlanExtend,
558    /// Skip tree search, direct token sampling (low uncertainty, confident).
559    PlanSkip,
560    /// Activate SpecHop continuous speculation with k speculative threads (Plan 131).
561    /// Selected when speculator latency α is low and tool ratio β is moderate.
562    SpecHop { k: usize },
563    /// Harness update: AbsorbCompress promote + HotSwapPruner reload (Plan 163 T5).
564    /// Selected when harness has plateaued and a compressed arm set may improve.
565    #[cfg(feature = "sia_feedback")]
566    HarnessUpdate,
567    /// Weight update: trigger riir-gpu training step on accumulated TrialLog (Plan 163 T6).
568    /// Selected when stall detection fires — reward plateau suggests weights need updating.
569    #[cfg(feature = "sia_feedback")]
570    WeightUpdate,
571}
572
573/// Context key for configurator bandit — coarse entropy binning.
574///
575/// Entropy is discretized into 10 bins via `floor(entropy * 10.0)` clamped to 0..9.
576/// Combined with domain index, this provides context-aware arm selection.
577#[cfg(feature = "sr2am_configurator")]
578#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
579pub struct ConfiguratorContext {
580    /// Domain index from bandit infrastructure.
581    pub domain: usize,
582    /// Coarse entropy bin: `floor(entropy * 10.0)`, clamped to 0..9.
583    /// u8 — values are 0..9, packed after usize to avoid padding.
584    pub entropy_bin: u8,
585    /// Coarse desperation bin: `floor(desperation * 10.0)`, clamped to 0..9.
586    /// Plan 162 T11: emotion vector desperation score as additional context.
587    /// 0 = not desperate, 9 = highly desperate.
588    pub desperation_bin: u8,
589    /// Coarse epiplexity bin: `floor(epiplexity * 10.0)`, clamped to 0..9.
590    /// Plan 130 T4: structural information content (S_T) as additional context.
591    /// 0 = no structure detectable, 9 = highly structured.
592    pub epiplexity_bin: u8,
593}
594
595#[cfg(feature = "sr2am_configurator")]
596impl ConfiguratorContext {
597    /// Create context without desperation information (legacy compatibility).
598    ///
599    /// Sets `desperation_bin` to 0 (not desperate). Use `with_desperation()`
600    /// when emotion vector data is available.
601    pub fn new(domain: usize, entropy_bin: usize) -> Self {
602        Self {
603            domain,
604            entropy_bin: (entropy_bin.min(9)) as u8,
605            desperation_bin: 0,
606            epiplexity_bin: 0,
607        }
608    }
609
610    /// Set the desperation bin from a raw desperation score.
611    ///
612    /// `floor(desperation * 10.0)`, clamped to 0..9.
613    pub fn with_desperation(mut self, desperation: f32) -> Self {
614        self.desperation_bin = ((desperation * 10.0).floor() as u8).min(9);
615        self
616    }
617
618    /// Set the epiplexity bin from a raw epiplexity score (S_T).
619    ///
620    /// `floor(epiplexity * 10.0)`, clamped to 0..9.
621    /// S_T measures structural information content — higher values indicate
622    /// more structure that a bounded observer can extract from the data.
623    pub fn with_epiplexity(mut self, epiplexity: f32) -> Self {
624        self.epiplexity_bin = ((epiplexity * 10.0).floor() as u8).min(9);
625        self
626    }
627
628    /// Create context from entropy and epiplexity signals.
629    ///
630    /// Convenience constructor that bins both entropy (H_T proxy) and
631    /// epiplexity (S_T structural information) in one call.
632    /// `desperation_bin` defaults to 0.
633    pub fn from_entropy_epiplexity(domain: usize, entropy: f32, epiplexity: f32) -> Self {
634        let entropy_bin = ((entropy * 10.0).floor() as u8).min(9);
635        let epiplexity_bin = ((epiplexity * 10.0).floor() as u8).min(9);
636        Self {
637            domain,
638            entropy_bin,
639            desperation_bin: 0,
640            epiplexity_bin,
641        }
642    }
643
644    /// Discretize epiplexity (S_T) into a coarse bin index.
645    ///
646    /// `floor(epiplexity * 10.0)`, clamped to 0..9.
647    pub fn epiplexity_bin(epiplexity: f32) -> u8 {
648        ((epiplexity * 10.0).floor() as u8).min(9)
649    }
650}
651
652// ---------------------------------------------------------------------------
653// EqR Convergence Selection (Plan 119)
654// ---------------------------------------------------------------------------
655
656/// Selection strategy for width-scaled rollouts (EqR convergence-based selection).
657///
658/// Maps to `WidthSelectionMode` (in `crate::speculative::dd_tree`) at runtime.
659/// This enum lives in `katgpt-core` so Config can reference it without depending on
660/// the speculative decode module.
661///
662/// - `BestQ`: Highest cumulative relevance (PTRM default, no behavior change)
663/// - `MajorityVote`: Most common path across rollouts (mode@K)
664/// - `Top1Converged`: Smallest final residual ∥p_{d+1} − p_d∥ (EqR proxy)
665/// - `BtRank`: Pairwise Bradley-Terry ranking (requires `bt_rank` feature)
666///
667/// **Precondition:** `Top1Converged` is only reliable after landscape shaping
668/// (RI + NI training). See Research 079 (EqR) for theoretical justification.
669#[repr(u8)]
670#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
671pub enum ConvergenceSelector {
672    /// Select rollout with highest cumulative relevance score (PTRM Q-head analog).
673    #[default]
674    BestQ,
675    /// Select the most frequent path across all rollouts (mode@K, majority vote).
676    MajorityVote,
677    /// Select rollout with smallest marginal-change residual ∥p_{d+1} − p_d∥ (EqR proxy).
678    Top1Converged,
679    /// Pairwise Bradley-Terry ranking across rollouts (requires `bt_rank` feature).
680    BtRank,
681}
682
683// ---------------------------------------------------------------------------
684// Wall Attention — Diagonal Forget Gates Replacing RoPE (Plan 173)
685// ---------------------------------------------------------------------------
686
687/// Wall Attention configuration (Plan 173, Research: Wall Attention paper).
688///
689/// Wall replaces RoPE with diagonal forget gates applied as factorized Q/K rescaling:
690/// `q̃_i = exp(P_i) ⊙ q_i`, `k̃_j = exp(-P_j) ⊙ k_j`.
691/// This means attention kernels are UNCHANGED — they receive pre-rescaled Q and K.
692///
693/// Only applicable to Wall-trained models (requires W_g gate projection weights).
694///
695/// The wall on/off switch lives at the parent `Config.wall_config: Option<WallConfig>`
696/// level — `None` means use RoPE/fallback, `Some(_)` means Wall is active. There is
697/// no `use_wall` field on the struct itself (the canonical design since Plan 173).
698#[derive(Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
699#[cfg(feature = "wall_attention")]
700pub struct WallConfig {
701    /// Gate bias initialization value. Default 6.0 = open gate (vanilla attention behavior).
702    /// Lower values → more active forgetting (gate_bias=0 → retention ≈ 0.62).
703    pub gate_bias: f32,
704    /// Maximum gate log-sigmoid clamp value. Default 0.87 (matches paper).
705    /// Gates are clamped to (-gate_max, 0] after log-sigmoid.
706    pub gate_max: f32,
707    /// Use key-projected gate variant (derive gate from K projection).
708    /// Preferred for zero KV cache overhead — gate is computed from key, not hidden state.
709    pub use_key_projected: bool,
710    /// Gate projection dimension = n_kv_heads * head_dim.
711    /// Default 0 = compute from model dims via [`Self::with_dims`] or
712    /// [`Self::validate`]. Set explicitly only when the consumer already knows
713    /// the dim and wants to skip the derivation.
714    pub gate_proj_dim: usize,
715}
716
717#[cfg(feature = "wall_attention")]
718impl Default for WallConfig {
719    fn default() -> Self {
720        Self {
721            gate_bias: 6.0,
722            gate_max: 0.87,
723            use_key_projected: true,
724            gate_proj_dim: 0,
725        }
726    }
727}
728
729#[cfg(feature = "wall_attention")]
730impl WallConfig {
731    pub fn new() -> Self {
732        Self::default()
733    }
734
735    /// Validate config consistency against model dimensions.
736    ///
737    /// Checks:
738    /// - `gate_max` is in the open interval (0, 1) — soft-clamp must produce
739    ///   finite non-degenerate log-gates.
740    /// - `gate_proj_dim` is either unset (0, meaning "derive from dims") or
741    ///   exactly `n_kv_heads * head_dim`.
742    pub fn validate(&self, n_kv_heads: usize, head_dim: usize) -> Result<(), String> {
743        if self.gate_max <= 0.0 || self.gate_max >= 1.0 {
744            return Err(format!("gate_max must be in (0, 1), got {}", self.gate_max));
745        }
746        let expected_dim = n_kv_heads * head_dim;
747        if self.gate_proj_dim != 0 && self.gate_proj_dim != expected_dim {
748            return Err(format!(
749                "gate_proj_dim ({}) must equal n_kv_heads * head_dim ({})",
750                self.gate_proj_dim, expected_dim
751            ));
752        }
753        Ok(())
754    }
755
756    /// Builder: derive `gate_proj_dim` from model dimensions.
757    /// Consumes and returns `self` for chaining.
758    pub fn with_dims(mut self, n_kv_heads: usize, head_dim: usize) -> Self {
759        self.gate_proj_dim = n_kv_heads * head_dim;
760        self
761    }
762}
763
764// ---------------------------------------------------------------------------
765// Collapse-Aware Adaptive Thinking (Plan 212)
766// ---------------------------------------------------------------------------
767
768/// Per-instance adaptive budget for collapse-aware thinking.
769///
770/// Controls when mid-reasoning early exit triggers and how efficiency rewards
771/// are shaped. Feature-gated behind `collapse_aware_thinking`.
772#[allow(dead_code)]
773#[derive(Clone, Copy, Debug)]
774pub struct ThinkingBudget {
775    /// Maximum thinking tokens before forced termination.
776    pub max_tokens: u32,
777    /// Hesitation count threshold τ — collapse triggers when exceeded.
778    pub collapse_threshold: u32,
779    /// Efficiency–accuracy trade-off for reward shaping.
780    /// Higher γ penalizes longer traces more aggressively.
781    /// Range: [0.0, 1.0].
782    pub efficiency_gamma: f32,
783}
784
785#[cfg(feature = "collapse_aware_thinking")]
786impl Default for ThinkingBudget {
787    fn default() -> Self {
788        Self {
789            max_tokens: 4096,
790            collapse_threshold: 3,
791            efficiency_gamma: 0.5,
792        }
793    }
794}