katgpt_types/
looping.rs

1//! Training-Free Loop types.
2
3use super::*;
4
5// ---------------------------------------------------------------------------
6// Training-Free Loop Types (Plan 136)
7// ---------------------------------------------------------------------------
8
9/// Sub-step integration strategy for the training-free loop.
10///
11/// Controls how intermediate loop outputs are combined with the running state.
12#[repr(u8)]
13#[derive(Clone, Copy, Debug, Default, PartialEq)]
14pub enum SubStepStrategy {
15    /// Damped Euler: x ← x + (1/K)·(y − x)
16    #[default]
17    DampedEuler,
18    /// K-stage Runge-Kutta: x ← β·y + (1−β)·x
19    KStageRK {
20        /// Blend factor β ∈ [0, 1]. 0.5 is neutral (equal weight).
21        beta: f32,
22    },
23}
24
25/// Iteration mode for the training-free loop window.
26///
27/// Controls whether the window is applied as a single block or iterated
28/// layer-by-layer within each sub-step.
29#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
30#[repr(u8)]
31pub enum IterationMode {
32    /// Apply the full window [a, b] as one block per sub-step.
33    #[default]
34    Block,
35    /// Apply each layer in the window individually per sub-step.
36    Layer,
37}
38
39/// KV cache write strategy for the training-free loop.
40///
41/// Controls which loop iteration writes the canonical KV entries.
42#[repr(u8)]
43#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
44pub enum CacheStrategy {
45    /// Use the final loop iteration's hidden state for KV cache.
46    #[default]
47    Last,
48    /// Use the pre-loop hidden state for KV cache (first iteration).
49    First,
50}
51
52/// Configuration for training-free loop wrapper (Plan 136).
53///
54/// Pure inference-time retrofit: re-applies a contiguous mid-stack block of
55/// layers with ODE-motivated damped sub-stepping. No training needed.
56#[derive(Clone, Debug)]
57pub struct TrainingFreeLoopConfig {
58    /// Start of the loop window (inclusive layer index).
59    pub window_start: usize,
60    /// End of the loop window (inclusive layer index).
61    pub window_end: usize,
62    /// Number of loop iterations (K in the paper).
63    pub loop_count: usize,
64    /// Sub-step integration strategy.
65    pub strategy: SubStepStrategy,
66    /// Iteration mode (block vs layer-wise).
67    pub iteration_mode: IterationMode,
68    /// KV cache write strategy.
69    pub cache_strategy: CacheStrategy,
70}
71
72impl Default for TrainingFreeLoopConfig {
73    fn default() -> Self {
74        Self {
75            window_start: 0,
76            window_end: 0,
77            loop_count: 2,
78            strategy: SubStepStrategy::KStageRK { beta: 0.5 },
79            iteration_mode: IterationMode::Block,
80            cache_strategy: CacheStrategy::First,
81        }
82    }
83}
84
85impl TrainingFreeLoopConfig {
86    /// Creates a config with sensible defaults for a given `Config`.
87    ///
88    /// Window heuristic: center at 48% depth, ±1 layer.
89    /// For small models (≤4 layers), defaults to (0, n_layer−1).
90    /// Uses the paper-recommended K-stage RK strategy with β=0.5.
91    pub fn from_config(config: &Config) -> Self {
92        let n = config.n_layer;
93        let (window_start, window_end) = if n <= 4 {
94            (0, n.saturating_sub(1))
95        } else {
96            let center = (n as f32 * 0.48) as usize;
97            (center.saturating_sub(1), (center + 2).min(n - 1))
98        };
99        Self {
100            window_start,
101            window_end,
102            loop_count: 2,
103            strategy: SubStepStrategy::KStageRK { beta: 0.5 },
104            iteration_mode: IterationMode::Block,
105            cache_strategy: CacheStrategy::First,
106        }
107    }
108}
katgpt_types/looping.rs

katgpt_types/
looping.rs
