katgpt-types 0.2.0

Shared configuration, RNG, math utilities, LoRA, domain embeddings, and inference types for katgpt-rs / riir-engine. Pure substrate leaf — no katgpt-* deps.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
//! Small configuration enums and feature-config structs.

// Shared configuration, RNG, and math utilities.
// Superset of types from both katgpt-rs and riir-engine projects.

// ---------------------------------------------------------------------------
// Enums
// ---------------------------------------------------------------------------

/// Adaptive depth tier mapping to layer count (Plan 284).
/// Reuses ThermalPath naming convention from FlashAR Consensus (Plan 166).
///
/// | Tier     | Layers     | When                              |
/// |----------|------------|-----------------------------------|
/// | Plasma   | 1          | High entropy, easy positions      |
/// | Hot      | 2          | Medium entropy, standard tactics  |
/// | Warm     | all        | Low entropy, complex positions    |
/// | Cold     | all+verify | Critical, full verification       |
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(u8)]
pub enum DepthTier {
    /// Easy positions: empty board, forced moves. 1 layer.
    Plasma = 0,
    /// Moderate: standard tactics. 2 layers.
    Hot = 1,
    /// Complex: all layers + spot-check verification.
    Warm = 2,
    /// Critical: all layers + full verification.
    Cold = 3,
}

impl DepthTier {
    /// Returns the maximum number of transformer layers to execute for this tier.
    pub fn max_layers(&self, total_layers: usize) -> usize {
        match self {
            DepthTier::Plasma => 1.min(total_layers),
            DepthTier::Hot => 2.min(total_layers),
            DepthTier::Warm => total_layers,
            DepthTier::Cold => total_layers,
        }
    }
}

/// Attention mode for HLA (Higher-order Linear Attention).
///
/// - `Standard`: SDPA with KV cache (default, backward-compatible).
/// - `Hla`: Symmetric second-order linear attention — O(1) per-token memory.
/// - `Ahla`: Asymmetric second-order linear attention — lower state cost than symmetric.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum HlaMode {
    #[default]
    Standard,
    /// Symmetric second-order: SK, CQV, mQ accumulators.
    Hla,
    /// Asymmetric second-order: PKV, mK accumulators.
    Ahla,
}

/// Attention mode for forward passes.
///
/// - `Causal`: Standard autoregressive — only attend to positions ≤ current (default).
/// - `Bidirectional`: Attend to ALL positions — used for dLLM masked prediction (Plan 066).
/// - `BlockCausal`: Bidirectional within current block, causal across blocks — D2F student.
/// - `SpKv`: SP-KV self-pruned key-value attention (Plan 070).
/// - `SpKvQuant`: SP-KV + Quantized KV fusion (Plan 070 Phase 3, Task T12).
#[repr(u8)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum AttentionMode {
    #[default]
    Causal,
    /// Full bidirectional: all positions see all positions (teacher mode).
    Bidirectional,
    /// Block-causal: bidirectional within block, causal across blocks (student mode).
    BlockCausal,
    /// SP-KV self-pruned key-value attention (Plan 070).
    /// Learns which KV pairs to retain via utility prediction.
    /// Gate bias = log(u) during training, 0|-inf during inference.
    SpKv,
    /// SP-KV + Quantized KV fusion (Plan 070 Phase 3, Task T12).
    /// Selective write (SP-KV utility gating) + lossy quantize (any QuantizedKVCache backend).
    /// Two-stage compression: only useful KV pairs kept, those compressed to 2-4 bits/coord.
    SpKvQuant,
    /// DashAttention: adaptive sparse hierarchical attention via α-entmax routing (Plan 106).
    /// Replaces fixed-budget top-k block selection with adaptive support selection.
    /// Learned chunk summaries via head_cls vectors.
    DashAttn,
}

/// Model architecture selector for forward pass dispatch.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum ModelArchitecture {
    #[default]
    Generic,
    Gemma2,
    Llama,
    /// Hybrid DeltaNet/Attention model (e.g., Qwen 3.5, Kimi Linear).
    /// Uses per-layer config to determine DeltaNet vs standard attention.
    /// Plan 182: Luce Megakernel Distill — DeltaNet GPU Inference.
    #[cfg(feature = "deltanet_inference")]
    QwenDeltaNet,
}

/// Attention projection configuration.
/// Controls whether K and V projections share weights (Q-K=V tying).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[repr(u8)]
pub enum AttentionProjection {
    /// Standard Q, K, V (3 projections, full KV cache)
    #[default]
    Full,
    /// Q-K=V: K and V share projection (2 projections, K-only cache).
    /// 50% KV cache reduction, ~3% perplexity cost.
    /// Post-hoc weight merging: W_kv = (W_k + W_v) / 2.
    SharedKV,
}

/// KV cache layout (derived from AttentionProjection).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum CacheLayout {
    /// Store both K and V (standard)
    KV,
    /// Store K only, V = K at read (SharedKV)
    K,
}

/// Weight storage dtype (affects loading and dequantization).
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum WeightDtype {
    #[default]
    F32,
    F16,
    BF16,
}

// ---------------------------------------------------------------------------
// Delta Routing (Plan 097, Research 061)
// ---------------------------------------------------------------------------

/// Delta routing mode — cross-layer information flow via delta vectors.
/// Research 061: Delta Attention Residuals (Plan 097).
///
/// Kept compiled even when `delta_routing` is off so config round-trips
/// serialize identically across feature sets. Reachable via `Config` defaults
/// once the routing backend lands.
#[allow(dead_code)]
#[repr(u8)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum DeltaRoutingMode {
    /// No delta routing (default).
    #[default]
    Off,
    /// Delta Block: accumulate deltas within blocks of `block_size` layers.
    /// B+1 sources per routing decision. ~20% throughput overhead.
    DeltaBlock,
    /// Delta Attention Residuals: per-sublayer delta routing.
    /// 2L sources. 69% throughput reduction at L=36. Use only for research.
    DeltaAttnRes,
}

/// Configuration for delta routing (Plan 097, Research 061).
///
/// Fields ordered by descending alignment to minimize padding:
/// usize (8B) → repr(u8) enum (1B) — 16 bytes total, no wasted padding.
#[allow(dead_code)]
#[derive(Clone, Copy, Debug)]
pub struct DeltaRoutingConfig {
    /// Block size for DeltaBlock mode (number of layers per block).
    /// Default: 4. Paper recommends B=4.
    pub block_size: usize,
    /// Routing mode.
    pub mode: DeltaRoutingMode,
}

impl Default for DeltaRoutingConfig {
    fn default() -> Self {
        Self {
            block_size: 4,
            mode: DeltaRoutingMode::Off,
        }
    }
}

// ---------------------------------------------------------------------------
// DeltaNet Inference (Plan 182: Luce Megakernel Distill)
// ---------------------------------------------------------------------------

/// Per-layer type for hybrid DeltaNet/Attention models.
/// Each layer is either a standard attention layer or a DeltaNet recurrent layer.
#[cfg(feature = "deltanet_inference")]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum DeltaNetLayerType {
    /// Standard multi-head attention with KV cache.
    #[default]
    Attention,
    /// DeltaNet linear recurrent layer (fast recurrent update, no KV cache needed).
    DeltaNet,
}

// DeltaRoutingConfig::delta_block / is_enabled are intended for the
// delta_routing backend (Plan 097) which is still being wired up. Silence
// dead-code until callers land.
#[allow(dead_code)]
impl DeltaRoutingConfig {
    pub fn delta_block(block_size: usize) -> Self {
        Self {
            mode: DeltaRoutingMode::DeltaBlock,
            block_size,
        }
    }

    pub fn is_enabled(&self) -> bool {
        self.mode != DeltaRoutingMode::Off
    }
}

// ---------------------------------------------------------------------------
// DashAttention Config (Plan 106, Research 68)
// ---------------------------------------------------------------------------

/// Configuration for DashAttention adaptive sparse hierarchical attention.
/// Controls α-entmax routing, chunk summarization, and routing bias.
///
/// Fields ordered by descending alignment to minimize padding:
/// usize (8B) → f32 (4B) → bool (1B) — 24 bytes total, no wasted padding.
#[derive(Clone, Copy, Debug)]
pub struct DashAttnConfig {
    /// Chunk size for block-level attention (default: 64).
    pub chunk_size: usize,
    /// α parameter for entmax. Only α=1.5 supported (quadratic, closed-form).
    pub alpha: f32,
    /// Scaling factor γ applied to chunk logits before entmax (default: 1.0).
    pub scaling_factor: f32,
    /// Prior strength σ for routing bias (default: 1e6, weak prior).
    pub sigma: f32,
    /// Whether to estimate diagonal attention contribution (default: true).
    /// Tail-packed after f32 group to avoid bool-between-f32 padding.
    pub estimate_diagonal: bool,
}

impl Default for DashAttnConfig {
    fn default() -> Self {
        Self {
            chunk_size: 64,
            alpha: 1.5,
            scaling_factor: 1.0,
            sigma: 1e6,
            estimate_diagonal: true,
        }
    }
}

// ---------------------------------------------------------------------------
// RTPurbo Retrieval Head Sparse Decode (Plan 126, Research 86)
// ---------------------------------------------------------------------------

/// Head role classification for RTPurbo sparse decode.
///
/// Only ~15% of attention heads ("retrieval heads") need full long-context access.
/// The remaining ~85% ("local heads") attend only to local context + attention sinks.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[repr(u8)]
pub enum RetrievalHeadRole {
    /// Local head — sliding window + sink tokens only, no full KV scan.
    #[default]
    Local,
    /// Retrieval head — low-dim projection + dynamic top-p token selection.
    Retrieval,
}

/// Configuration for RTPurbo retrieval head sparse decode.
///
/// Feature gate: `rt_turbo` (opt-in, requires `dash_attn`).
/// Adds head-wise retrieval/local classification + dynamic top-p token selection
/// for decode-phase sparse attention. Complements DashAttention's α-entmax block
/// routing with per-head specialization.
///
/// Must pass 6/6 GOAT proofs before default-on promotion.
///
/// Fields ordered by descending alignment to minimize padding:
/// usize (8B) → f32 (4B) → CalibrationMode (1B) — no padding between groups.
///
/// # Calibration mode (Plan 358)
///
/// [`CalibrationMode::AttentionMass`] is the default (cheaper: 1 forward pass).
/// [`CalibrationMode::CausalNecessity`] is opt-in — strictly stronger on
/// workloads with correlated bystanders but ~10–100× more expensive to
/// calibrate. See `calibrate_from_causal_scores` in `rt_turbo::calibration`.
#[derive(Clone, Copy, Debug)]
pub struct RtTurboConfig {
    /// Low-dimensional projection size for pre-RoPE scoring (default: 16).
    /// Paper ablation: dim=16 is the sweet spot for low-frequency retrieval.
    pub low_dim: usize,
    /// Sliding window size for local heads (default: 8192).
    pub sliding_window: usize,
    /// Number of attention sink tokens always retained for local heads (default: 4).
    pub sink_tokens: usize,
    /// Block size for block-level top-p variant (default: 64).
    /// Should match `DashAttnConfig::chunk_size` for consistent routing.
    pub block_size: usize,
    /// Fraction of heads classified as retrieval heads (default: 0.15).
    /// Paper ablation: 15% is optimal balance of accuracy vs sparsity.
    pub retrieval_head_ratio: f32,
    /// Cumulative attention mass threshold for dynamic top-p selection (default: 0.9).
    /// Paper ablation: top-p=0.9 preserves >93% attention mass at 97% sparsity.
    pub top_p: f32,
    /// Which score semantics to use for head calibration (Plan 358). Default:
    /// `AttentionMass` (cheaper). `CausalNecessity` is opt-in — strictly
    /// stronger on bystander-heavy workloads but ~10–100× more expensive.
    /// `AdaptiveCausal` (Proposal 004) is opt-in — cheap-proxy escalate,
    /// unvalidated, requires per-head OV norms from the caller.
    pub calibration_mode: CalibrationMode,
}

/// Head-calibration score source (Plan 358, Research 362).
///
/// `#[repr(u8)]` for sync-friendly 1-byte representation.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
#[repr(u8)]
pub enum CalibrationMode {
    /// Observational needle attention-mass (RTPurbo Plan 126 default).
    /// Cheaper: a single forward pass + per-head mass scan.
    #[default]
    AttentionMass = 0,
    /// Causal necessity via activation/path patching IE score (Plan 358).
    /// Strictly stronger — excludes correlated bystanders — but requires
    /// `n_heads × n_calibration_samples` patched forward passes. Requires the
    /// `causal_head_importance` feature on the consuming crate.
    CausalNecessity = 1,
    /// Adaptive cheap-proxy escalate (Proposal 004 — OUR INVENTION, not from
    /// HydraHead). Uses an OV-circuit proxy (`attention_mass / ||OV_out||`)
    /// to detect bystander suspects, then escalates to Plan 358's causal
    /// patching only on those `k` suspects instead of all `n_heads`. Pays zero
    /// patched forwards when there are no bystanders (degenerates to
    /// `AttentionMass`). Requires the `adaptive_causal_calibration` feature.
    ///
    /// **UNVALIDATED.** Promotion to default is blocked on G1 (proxy precision)
    /// and G2 (cost reduction), both deferred to riir-engine. Unlike the other
    /// two modes, the caller must supply per-head OV output norms (from a real
    /// transformer forward) — see `calibrate_from_adaptive_causal`.
    AdaptiveCausal = 2,
}

impl Default for RtTurboConfig {
    fn default() -> Self {
        Self {
            low_dim: 16,
            sliding_window: 8192,
            sink_tokens: 4,
            block_size: 64,
            retrieval_head_ratio: 0.15,
            top_p: 0.9,
            calibration_mode: CalibrationMode::default(),
        }
    }
}

// ---------------------------------------------------------------------------
// LT2 Looped Inference (Plan 108, Research 73)
// ---------------------------------------------------------------------------

/// Looped transformer mode — weight-shared layer repetition.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum LoopMode {
    /// Standard single-pass (no looping).
    #[default]
    None,
    /// Weight-shared looping: same layers applied T times.
    /// Effective depth = n_layer × loop_count.
    WeightShared { loop_count: usize },
    /// Training-free loop: ODE-refined sub-stepping over a window of layers.
    /// No extra parameters — pure inference-time retrofit (Plan 136).
    TrainingFree,
}

/// Hybrid attention pattern for looped inference.
/// Controls which layers use full SDPA vs linear attention.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum HybridPattern {
    /// All layers use the same attention mode.
    #[default]
    Uniform,
    /// Depth-level interleave: every Nth layer uses full SDPA.
    /// e.g., Interleave { full_ratio: 5 } = every 5th layer is full.
    /// Paper optimal: 1:4 ratio (full_ratio=5).
    Interleave { full_ratio: usize },
    /// Bookend: first and last layers are full, middle is linear.
    Bookend,
}

/// Loop stability mode for weight-shared looped inference (Plan 428).
///
/// Parameter-free architectural fixes for T-pass loop stability, validated
/// via the §3.6 defend-wrong PoC benchmark.
///
/// **Only `InterLoopNorm` ships** — the PoC proved it's the sole fix that
/// controls residual norm growth. FLA-res (direct residual addition of
/// `prev_h` at every layer) caused catastrophic norm explosion (~2.2B× at
/// T=12), and Attention Injection was a no-op for single-position attention
/// (softmax of 1 element = 1.0, so Q doesn't affect the output). Both were
/// dropped per the defend-wrong verdict.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[repr(u8)]
pub enum LoopStabilityMode {
    /// No inter-loop stabilization (byte-identical to pre-Plan-428 behavior).
    #[default]
    None,
    /// Inter-loop RMSNorm: normalize the hidden state between loop iterations
    /// (tau > 0), before the inner layer pass. PoC: norm ratio 3.34× vs
    /// baseline 11.19×, KL 0.0008, step-size trend converging (14.9 → 2.05).
    InterLoopNorm,
}

/// Head-specific sigmoid gate after SDPA, before Wo.
/// Zero-initialized → starts at sigmoid(0) = 0.5 (neutral multiplicative identity).
#[derive(Clone)]
pub struct SdpaOutputGate {
    /// Gate weights: [n_heads * head_dim, dim].
    /// Zero-init so gate starts at sigmoid(0) = 0.5.
    pub w_gate: Vec<f32>,
}

impl SdpaOutputGate {
    /// Allocate zeroed gate weights.
    pub fn new(n_heads: usize, head_dim: usize, dim: usize) -> Self {
        Self {
            w_gate: vec![0.0; n_heads * head_dim * dim],
        }
    }

    /// Apply sigmoid-gated projection to attention output.
    ///
    /// Computes: `gate[i] = sigmoid(W_gate[i] · attn_out)`, then `attn_out[i] *= gate[i]`.
    /// Zero-init weights produce sigmoid(0) = 0.5 for all (neutral half-pass).
    /// Paper reference: +0.3–0.5 avg points on zero-shot benchmarks.
    pub fn forward(&self, attn_out: &mut [f32], dim: usize, temp: &mut [f32]) {
        let n = attn_out.len();
        debug_assert!(temp.len() >= n, "temp buffer too small");
        debug_assert!(self.w_gate.len() >= n * dim, "gate weights too small");

        // Step 1: Compute gate signal = sigmoid(W_gate @ attn_out)
        // Batch matvec then batch sigmoid avoids per-element loop overhead
        crate::simd::simd_matvec(temp, &self.w_gate, attn_out, n, dim);

        // SIMD sigmoid: temp = -temp, exp, then 1/(1+exp)
        crate::simd::simd_scale_inplace(&mut temp[..n], -1.0);
        crate::simd::simd_exp_inplace(&mut temp[..n]);
        crate::simd::simd_add_scalar_inplace(&mut temp[..n], 1.0);
        // temp now = 1 + exp(-x), invert: temp = 1/temp = sigmoid
        crate::simd::simd_reciprocal_inplace(&mut temp[..n]);

        // Step 2: Apply gate elementwise via SIMD scale-mul (fused)
        // attn_out[i] *= temp[i] is element-wise multiply
        // Use simd_scale_mul_inplace with scale=1.0: attn[i] = temp[i] * attn[i] * 1.0
        crate::simd::simd_scale_mul_inplace(attn_out, &temp[..n], 1.0);
    }
}

/// Per-loop residual scaling gate.
/// h^(τ) = h̃^(τ) + ρ_τ ⊙ h^(τ-1)
/// Zero-init so first iteration is h̃^(1) (no residual from "previous").
#[derive(Clone)]
pub struct ResidualGate {
    /// Per-loop gates: [loop_count, dim].
    /// Each ρ_τ is element-wise, zero-init.
    pub gates: Vec<f32>,
}

impl ResidualGate {
    /// Allocate zeroed residual gates.
    pub fn new(loop_count: usize, dim: usize) -> Self {
        Self {
            gates: vec![0.0; loop_count * dim],
        }
    }

    /// Deterministic loop-stable residual gates (Plan 483 T2.1, §3.5 path 2).
    ///
    /// Sets gates to a constant `decay` factor for τ > 0, enabling information
    /// carry-forward between T passes. The first loop (τ=0) has zero gates
    /// (no previous state to carry forward).
    ///
    /// This is a **modelless** construction — no training, no gradient descent.
    /// The `decay` factor controls the trade-off between carry-forward strength
    /// and stability. Conservative values (0.1–0.3) are safe for most weight
    /// matrices; larger values risk divergence.
    ///
    /// Rationale: the zero-init default (`new()`) makes every T-pass effectively
    /// independent — no hidden state carries forward between loops. This
    /// undermines the LT2 paper's "effective depth T×n_layer" claim. A non-zero
    /// deterministic gate restores the residual connection across loops without
    /// requiring trained gate parameters.
    ///
    /// # Arguments
    /// * `loop_count` - Number of T-passes (T)
    /// * `dim` - Hidden dimension (n_embd)
    /// * `decay` - Constant gate value for τ > 0 (e.g., 0.1 = conservative)
    #[inline]
    pub fn new_loop_stable(loop_count: usize, dim: usize, decay: f32) -> Self {
        let mut gates = vec![0.0f32; loop_count * dim];
        // τ=0: zero (no previous state). τ>0: constant decay.
        for tau in 1..loop_count {
            let offset = tau * dim;
            gates[offset..offset + dim].fill(decay);
        }
        Self { gates }
    }

    /// Deterministic loop-stable residual gates with exponential decay
    /// (Plan 483 T2.1, §3.5 path 2 variant).
    ///
    /// ρ_τ = `base`^(τ-1) for τ > 0 — later loops contribute exponentially less.
    /// This schedule mirrors the spectral-radius-based stabilization where the
    /// residual contribution decays as the hidden state converges.
    ///
    /// # Arguments
    /// * `loop_count` - Number of T-passes (T)
    /// * `dim` - Hidden dimension (n_embd)
    /// * `base` - Decay base (e.g., 0.5 → ρ_1=1.0, ρ_2=0.5, ρ_3=0.25, ...)
    #[inline]
    pub fn new_loop_stable_exp_decay(loop_count: usize, dim: usize, base: f32) -> Self {
        let mut gates = vec![0.0f32; loop_count * dim];
        for tau in 1..loop_count {
            let offset = tau * dim;
            let val = base.powi((tau - 1) as i32);
            gates[offset..offset + dim].fill(val);
        }
        Self { gates }
    }
}

// ---------------------------------------------------------------------------
// SR²AM Configurator Bandit (Plan 112, Research 076)
// ---------------------------------------------------------------------------

/// SR²AM Configurator decision — learned per-turn planning regulation.
///
/// The configurator selects one of these arms per inference turn based on
/// context (domain + entropy bin). UCB1 balances exploration vs exploitation.
///
/// - `PlanNew`: reset tree, full budget allocation (high uncertainty, new sub-problem)
/// - `PlanExtend`: keep tree, extend depth by one level (moderate uncertainty, continuing)
/// - `PlanSkip`: skip tree search, direct token sampling (low uncertainty, confident)
#[cfg(feature = "sr2am_configurator")]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
#[repr(u8)]
pub enum PlanningDecision {
    /// Reset tree, full budget allocation (high uncertainty, new sub-problem).
    PlanNew,
    /// Keep tree, extend depth by one level (moderate uncertainty, continuing).
    PlanExtend,
    /// Skip tree search, direct token sampling (low uncertainty, confident).
    PlanSkip,
    /// Activate SpecHop continuous speculation with k speculative threads (Plan 131).
    /// Selected when speculator latency α is low and tool ratio β is moderate.
    SpecHop { k: usize },
    /// Harness update: AbsorbCompress promote + HotSwapPruner reload (Plan 163 T5).
    /// Selected when harness has plateaued and a compressed arm set may improve.
    #[cfg(feature = "sia_feedback")]
    HarnessUpdate,
    /// Weight update: trigger riir-gpu training step on accumulated TrialLog (Plan 163 T6).
    /// Selected when stall detection fires — reward plateau suggests weights need updating.
    #[cfg(feature = "sia_feedback")]
    WeightUpdate,
}

/// Context key for configurator bandit — coarse entropy binning.
///
/// Entropy is discretized into 10 bins via `floor(entropy * 10.0)` clamped to 0..9.
/// Combined with domain index, this provides context-aware arm selection.
#[cfg(feature = "sr2am_configurator")]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ConfiguratorContext {
    /// Domain index from bandit infrastructure.
    pub domain: usize,
    /// Coarse entropy bin: `floor(entropy * 10.0)`, clamped to 0..9.
    /// u8 — values are 0..9, packed after usize to avoid padding.
    pub entropy_bin: u8,
    /// Coarse desperation bin: `floor(desperation * 10.0)`, clamped to 0..9.
    /// Plan 162 T11: emotion vector desperation score as additional context.
    /// 0 = not desperate, 9 = highly desperate.
    pub desperation_bin: u8,
    /// Coarse epiplexity bin: `floor(epiplexity * 10.0)`, clamped to 0..9.
    /// Plan 130 T4: structural information content (S_T) as additional context.
    /// 0 = no structure detectable, 9 = highly structured.
    pub epiplexity_bin: u8,
}

#[cfg(feature = "sr2am_configurator")]
impl ConfiguratorContext {
    /// Create context without desperation information (legacy compatibility).
    ///
    /// Sets `desperation_bin` to 0 (not desperate). Use `with_desperation()`
    /// when emotion vector data is available.
    pub fn new(domain: usize, entropy_bin: usize) -> Self {
        Self {
            domain,
            entropy_bin: (entropy_bin.min(9)) as u8,
            desperation_bin: 0,
            epiplexity_bin: 0,
        }
    }

    /// Set the desperation bin from a raw desperation score.
    ///
    /// `floor(desperation * 10.0)`, clamped to 0..9.
    pub fn with_desperation(mut self, desperation: f32) -> Self {
        self.desperation_bin = ((desperation * 10.0).floor() as u8).min(9);
        self
    }

    /// Set the epiplexity bin from a raw epiplexity score (S_T).
    ///
    /// `floor(epiplexity * 10.0)`, clamped to 0..9.
    /// S_T measures structural information content — higher values indicate
    /// more structure that a bounded observer can extract from the data.
    pub fn with_epiplexity(mut self, epiplexity: f32) -> Self {
        self.epiplexity_bin = ((epiplexity * 10.0).floor() as u8).min(9);
        self
    }

    /// Create context from entropy and epiplexity signals.
    ///
    /// Convenience constructor that bins both entropy (H_T proxy) and
    /// epiplexity (S_T structural information) in one call.
    /// `desperation_bin` defaults to 0.
    pub fn from_entropy_epiplexity(domain: usize, entropy: f32, epiplexity: f32) -> Self {
        let entropy_bin = ((entropy * 10.0).floor() as u8).min(9);
        let epiplexity_bin = ((epiplexity * 10.0).floor() as u8).min(9);
        Self {
            domain,
            entropy_bin,
            desperation_bin: 0,
            epiplexity_bin,
        }
    }

    /// Discretize epiplexity (S_T) into a coarse bin index.
    ///
    /// `floor(epiplexity * 10.0)`, clamped to 0..9.
    pub fn epiplexity_bin(epiplexity: f32) -> u8 {
        ((epiplexity * 10.0).floor() as u8).min(9)
    }
}

// ---------------------------------------------------------------------------
// EqR Convergence Selection (Plan 119)
// ---------------------------------------------------------------------------

/// Selection strategy for width-scaled rollouts (EqR convergence-based selection).
///
/// Maps to `WidthSelectionMode` (in `crate::speculative::dd_tree`) at runtime.
/// This enum lives in `katgpt-core` so Config can reference it without depending on
/// the speculative decode module.
///
/// - `BestQ`: Highest cumulative relevance (PTRM default, no behavior change)
/// - `MajorityVote`: Most common path across rollouts (mode@K)
/// - `Top1Converged`: Smallest final residual ∥p_{d+1} − p_d∥ (EqR proxy)
/// - `BtRank`: Pairwise Bradley-Terry ranking (requires `bt_rank` feature)
///
/// **Precondition:** `Top1Converged` is only reliable after landscape shaping
/// (RI + NI training). See Research 079 (EqR) for theoretical justification.
#[repr(u8)]
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ConvergenceSelector {
    /// Select rollout with highest cumulative relevance score (PTRM Q-head analog).
    #[default]
    BestQ,
    /// Select the most frequent path across all rollouts (mode@K, majority vote).
    MajorityVote,
    /// Select rollout with smallest marginal-change residual ∥p_{d+1} − p_d∥ (EqR proxy).
    Top1Converged,
    /// Pairwise Bradley-Terry ranking across rollouts (requires `bt_rank` feature).
    BtRank,
}

// ---------------------------------------------------------------------------
// Wall Attention — Diagonal Forget Gates Replacing RoPE (Plan 173)
// ---------------------------------------------------------------------------

/// Wall Attention configuration (Plan 173, Research: Wall Attention paper).
///
/// Wall replaces RoPE with diagonal forget gates applied as factorized Q/K rescaling:
/// `q̃_i = exp(P_i) ⊙ q_i`, `k̃_j = exp(-P_j) ⊙ k_j`.
/// This means attention kernels are UNCHANGED — they receive pre-rescaled Q and K.
///
/// Only applicable to Wall-trained models (requires W_g gate projection weights).
///
/// The wall on/off switch lives at the parent `Config.wall_config: Option<WallConfig>`
/// level — `None` means use RoPE/fallback, `Some(_)` means Wall is active. There is
/// no `use_wall` field on the struct itself (the canonical design since Plan 173).
#[derive(Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
#[cfg(feature = "wall_attention")]
pub struct WallConfig {
    /// Gate bias initialization value. Default 6.0 = open gate (vanilla attention behavior).
    /// Lower values → more active forgetting (gate_bias=0 → retention ≈ 0.62).
    pub gate_bias: f32,
    /// Maximum gate log-sigmoid clamp value. Default 0.87 (matches paper).
    /// Gates are clamped to (-gate_max, 0] after log-sigmoid.
    pub gate_max: f32,
    /// Use key-projected gate variant (derive gate from K projection).
    /// Preferred for zero KV cache overhead — gate is computed from key, not hidden state.
    pub use_key_projected: bool,
    /// Gate projection dimension = n_kv_heads * head_dim.
    /// Default 0 = compute from model dims via [`Self::with_dims`] or
    /// [`Self::validate`]. Set explicitly only when the consumer already knows
    /// the dim and wants to skip the derivation.
    pub gate_proj_dim: usize,
}

#[cfg(feature = "wall_attention")]
impl Default for WallConfig {
    fn default() -> Self {
        Self {
            gate_bias: 6.0,
            gate_max: 0.87,
            use_key_projected: true,
            gate_proj_dim: 0,
        }
    }
}

#[cfg(feature = "wall_attention")]
impl WallConfig {
    pub fn new() -> Self {
        Self::default()
    }

    /// Validate config consistency against model dimensions.
    ///
    /// Checks:
    /// - `gate_max` is in the open interval (0, 1) — soft-clamp must produce
    ///   finite non-degenerate log-gates.
    /// - `gate_proj_dim` is either unset (0, meaning "derive from dims") or
    ///   exactly `n_kv_heads * head_dim`.
    pub fn validate(&self, n_kv_heads: usize, head_dim: usize) -> Result<(), String> {
        if self.gate_max <= 0.0 || self.gate_max >= 1.0 {
            return Err(format!("gate_max must be in (0, 1), got {}", self.gate_max));
        }
        let expected_dim = n_kv_heads * head_dim;
        if self.gate_proj_dim != 0 && self.gate_proj_dim != expected_dim {
            return Err(format!(
                "gate_proj_dim ({}) must equal n_kv_heads * head_dim ({})",
                self.gate_proj_dim, expected_dim
            ));
        }
        Ok(())
    }

    /// Builder: derive `gate_proj_dim` from model dimensions.
    /// Consumes and returns `self` for chaining.
    pub fn with_dims(mut self, n_kv_heads: usize, head_dim: usize) -> Self {
        self.gate_proj_dim = n_kv_heads * head_dim;
        self
    }
}

// ---------------------------------------------------------------------------
// Collapse-Aware Adaptive Thinking (Plan 212)
// ---------------------------------------------------------------------------

/// Per-instance adaptive budget for collapse-aware thinking.
///
/// Controls when mid-reasoning early exit triggers and how efficiency rewards
/// are shaped. Feature-gated behind `collapse_aware_thinking`.
#[allow(dead_code)]
#[derive(Clone, Copy, Debug)]
pub struct ThinkingBudget {
    /// Maximum thinking tokens before forced termination.
    pub max_tokens: u32,
    /// Hesitation count threshold τ — collapse triggers when exceeded.
    pub collapse_threshold: u32,
    /// Efficiency–accuracy trade-off for reward shaping.
    /// Higher γ penalizes longer traces more aggressively.
    /// Range: [0.0, 1.0].
    pub efficiency_gamma: f32,
}

#[cfg(feature = "collapse_aware_thinking")]
impl Default for ThinkingBudget {
    fn default() -> Self {
        Self {
            max_tokens: 4096,
            collapse_threshold: 3,
            efficiency_gamma: 0.5,
        }
    }
}