feat(batch10): attention sinks (StreamingLLM), chunked prefill, per-layer LR decay

- Attention sinks (arXiv:2309.17453): AttentionSinkEviction always retains first
  sink_size KV positions + last window_size; evicts middle band in O(evict_count);
  select_evict_positions/should_retain consistent; 15 tests
- Chunked prefill (vLLM arXiv:2309.06180): ChunkedPrefillScheduler splits long
  prompts into chunk_size=512 chunks interleaved with decode steps (max 128 decode
  tokens/step); PrefillChunkState tracks progress/remaining/completion; drain_completed();
  14 tests including 1500-token→3-chunk coverage
- Per-layer LR decay (ULMFiT / discriminative fine-tuning): LayerLrDecayConfig with
  base_lr * decay_rate^(num_layers-1-depth); LayerLrDecayBuilder parses layer/layers/
  h/blocks/bracket notation param names; LayerLrScheduler with outer multiplier for
  cosine/linear schedule composition; 14 tests + 1 doctest

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

crates/production/rtx-inference/src/cache/attention_sink.rs

@@ -0,0 +1,471 @@
+//! StreamingLLM-style attention sink eviction.
+//!
+//! In StreamingLLM (Xiao et al., arXiv:2309.17453) the first few tokens in a
+//! sequence accumulate disproportionately large attention weights regardless of
+//! their semantic content — these positions are called **attention sinks**. By
+//! permanently retaining the first `sink_size` KV positions together with a
+//! sliding window of `window_size` recent positions, a model can run inference
+//! on sequences of unbounded length without memory growth.
+//!
+//! # Eviction semantics
+//!
+//! ```text
+//! [ 0 .. sink_size )          — always retained (sink tokens)
+//! [ sink_size .. n-window_size) — evicted when total > max_capacity()
+//! [ n-window_size .. n )      — always retained (recent working context)
+//! ```
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_inference::AttentionSinkEviction;
+//!
+//! let eviction = AttentionSinkEviction::new(4, 32);
+//! assert_eq!(eviction.max_capacity(), 36);
+//!
+//! // Under capacity — nothing to evict.
+//! assert!(eviction.select_evict_positions(36).is_empty());
+//!
+//! // One over capacity — exactly one middle position is evicted.
+//! let to_evict = eviction.select_evict_positions(37);
+//! assert_eq!(to_evict, vec![4]);
+//! ```
+
+/// StreamingLLM-style attention sink eviction.
+///
+/// Always retains the first `sink_size` KV positions (sink tokens) and the
+/// most recent `window_size` positions (working context). All positions in the
+/// middle are evicted when the sequence exceeds `sink_size + window_size`.
+///
+/// # Reference
+///
+/// StreamingLLM (Xiao et al., arXiv:2309.17453): efficient streaming inference
+/// via attention sinks, enabling infinite context without memory growth.
+///
+/// # Example
+///
+/// ```rust
+/// use rtx_inference::AttentionSinkEviction;
+///
+/// let sink = AttentionSinkEviction::new(4, 32);
+///
+/// // Positions 0..4 and the last 32 positions are always kept.
+/// assert!(sink.should_retain(0, 100));
+/// assert!(sink.should_retain(3, 100));
+/// assert!(sink.should_retain(99, 100));
+/// assert!(!sink.should_retain(4, 100));
+/// ```
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct AttentionSinkEviction {
+    /// Number of leading positions treated as attention sinks (always kept).
+    pub sink_size: usize,
+    /// Number of trailing positions kept as the sliding working context.
+    pub window_size: usize,
+}
+
+impl AttentionSinkEviction {
+    /// Create a new `AttentionSinkEviction`.
+    ///
+    /// Both `sink_size` and `window_size` may be zero; a zero value for either
+    /// means that component contributes no protected positions.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_inference::AttentionSinkEviction;
+    ///
+    /// let sink = AttentionSinkEviction::new(4, 32);
+    /// assert_eq!(sink.sink_size, 4);
+    /// assert_eq!(sink.window_size, 32);
+    /// ```
+    pub fn new(sink_size: usize, window_size: usize) -> Self {
+        Self {
+            sink_size,
+            window_size,
+        }
+    }
+
+    /// Returns the total capacity at which eviction first activates.
+    ///
+    /// Eviction occurs only when `total_positions > max_capacity()`. Sequences
+    /// at or below this length are stored without modification.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_inference::AttentionSinkEviction;
+    ///
+    /// let sink = AttentionSinkEviction::new(4, 32);
+    /// assert_eq!(sink.max_capacity(), 36);
+    /// ```
+    #[inline]
+    pub fn max_capacity(&self) -> usize {
+        self.sink_size.saturating_add(self.window_size)
+    }
+
+    /// Returns the number of positions that would be evicted for a sequence of
+    /// `total_positions` tokens.
+    ///
+    /// This equals `max(0, total_positions − max_capacity())`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_inference::AttentionSinkEviction;
+    ///
+    /// let sink = AttentionSinkEviction::new(4, 32);
+    /// assert_eq!(sink.evict_count(36), 0);
+    /// assert_eq!(sink.evict_count(37), 1);
+    /// assert_eq!(sink.evict_count(40), 4);
+    /// ```
+    #[inline]
+    pub fn evict_count(&self, total_positions: usize) -> usize {
+        total_positions.saturating_sub(self.max_capacity())
+    }
+
+    /// Selects the position indices to evict for a sequence of `total_positions`
+    /// tokens currently in the cache.
+    ///
+    /// Returns a **sorted, ascending** `Vec<usize>` of position indices to drop.
+    /// Returns an empty `Vec` when `total_positions <= max_capacity()`.
+    ///
+    /// The evicted range is always the contiguous middle band
+    /// `sink_size .. (total_positions − window_size)`, which are the positions
+    /// that are neither sink tokens nor part of the recent sliding window.
+    ///
+    /// # Invariants
+    ///
+    /// - Positions `0 .. sink_size` are **never** in the returned set.
+    /// - Positions `(total_positions − window_size) .. total_positions` are
+    ///   **never** in the returned set.
+    /// - The returned `Vec` is sorted in ascending order.
+    /// - If `sink_size + window_size >= total_positions`, returns `[]`.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_inference::AttentionSinkEviction;
+    ///
+    /// let sink = AttentionSinkEviction::new(2, 3);
+    /// // 7 positions, capacity 5 → evict exactly positions [2, 3].
+    /// assert_eq!(sink.select_evict_positions(7), vec![2usize, 3]);
+    /// ```
+    pub fn select_evict_positions(&self, total_positions: usize) -> Vec<usize> {
+        let count = self.evict_count(total_positions);
+        if count == 0 {
+            return Vec::new();
+        }
+
+        // The evictable range starts just after the sinks and ends just before
+        // the recent window. The range is always non-empty here because
+        // `count > 0` implies `total_positions > sink_size + window_size`, so
+        // `evict_start < evict_end`.
+        let evict_start = self.sink_size;
+        let evict_end = total_positions.saturating_sub(self.window_size);
+
+        // Both halves (sinks + window) together are fewer than total_positions,
+        // so the range [evict_start, evict_end) is well-defined and non-empty.
+        (evict_start..evict_end).collect()
+    }
+
+    /// Returns `true` when position `pos` should be retained in the cache given
+    /// that the sequence currently has `total_positions` tokens.
+    ///
+    /// A position is retained if it falls in the sink region `[0, sink_size)` OR
+    /// in the recent-window region `[total_positions − window_size, total_positions)`.
+    ///
+    /// When `total_positions <= max_capacity()` every position is retained.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use rtx_inference::AttentionSinkEviction;
+    ///
+    /// let sink = AttentionSinkEviction::new(4, 32);
+    /// // Sink region.
+    /// assert!(sink.should_retain(0, 100));
+    /// assert!(sink.should_retain(3, 100));
+    /// // Recent window.
+    /// assert!(sink.should_retain(68, 100));
+    /// assert!(sink.should_retain(99, 100));
+    /// // Middle position evicted.
+    /// assert!(!sink.should_retain(4, 100));
+    /// assert!(!sink.should_retain(67, 100));
+    /// ```
+    #[inline]
+    pub fn should_retain(&self, pos: usize, total_positions: usize) -> bool {
+        // Nothing to evict — keep everything.
+        if total_positions <= self.max_capacity() {
+            return true;
+        }
+
+        // Sink region.
+        if pos < self.sink_size {
+            return true;
+        }
+
+        // Recent window region.
+        let window_start = total_positions.saturating_sub(self.window_size);
+        pos >= window_start
+    }
+}
+
+impl Default for AttentionSinkEviction {
+    /// Default: 4 sink tokens, 128-token sliding window (reasonable for LLMs).
+    fn default() -> Self {
+        Self::new(4, 128)
+    }
+}
+
+// ── Tests ─────────────────────────────────────────────────────────────────────
+
+#[cfg(test)]
+mod attention_sink_tests {
+    use super::AttentionSinkEviction;
+
+    // ── select_evict_positions / evict_count ─────────────────────────────────
+
+    #[test]
+    fn test_no_eviction_when_under_capacity() {
+        let sink = AttentionSinkEviction::new(4, 32);
+        // Several values below max_capacity (36) must yield an empty vec.
+        for n in 0..36 {
+            let evicted = sink.select_evict_positions(n);
+            assert!(
+                evicted.is_empty(),
+                "expected empty for n={n}, got {evicted:?}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_no_eviction_at_exact_capacity() {
+        let sink = AttentionSinkEviction::new(4, 32);
+        assert_eq!(sink.max_capacity(), 36);
+        assert!(
+            sink.select_evict_positions(36).is_empty(),
+            "at exact capacity nothing should be evicted"
+        );
+    }
+
+    #[test]
+    fn test_single_eviction_one_over() {
+        let sink = AttentionSinkEviction::new(4, 32);
+        // 37 positions, capacity 36 → exactly 1 middle position (index 4) evicted.
+        let evicted = sink.select_evict_positions(37);
+        assert_eq!(evicted, vec![4usize], "one over capacity → evict exactly index 4");
+    }
+
+    #[test]
+    fn test_evicts_middle_positions_only() {
+        let sink = AttentionSinkEviction::new(4, 8);
+        // capacity = 12; use 20 positions → 8 positions evicted (indices 4..12).
+        let evicted = sink.select_evict_positions(20);
+        assert_eq!(evicted.len(), 8);
+
+        // Sink positions 0..4 must never appear.
+        for pos in 0..4usize {
+            assert!(!evicted.contains(&pos), "sink position {pos} must never be evicted");
+        }
+
+        // Recent window positions 12..20 must never appear.
+        for pos in 12..20usize {
+            assert!(!evicted.contains(&pos), "window position {pos} must never be evicted");
+        }
+
+        // Every evicted index must be in the middle band [4, 12).
+        for &pos in &evicted {
+            assert!(
+                (4..12).contains(&pos),
+                "evicted position {pos} is outside the expected middle band [4,12)"
+            );
+        }
+    }
+
+    #[test]
+    fn test_evict_count_formula() {
+        let sink = AttentionSinkEviction::new(4, 32);
+        let cap = sink.max_capacity();
+
+        // Below capacity: always 0.
+        for n in 0..=cap {
+            assert_eq!(
+                sink.evict_count(n),
+                0,
+                "evict_count({n}) should be 0 when n <= cap"
+            );
+        }
+
+        // Above capacity: n - cap.
+        for extra in 1..=100 {
+            let n = cap + extra;
+            assert_eq!(
+                sink.evict_count(n),
+                extra,
+                "evict_count({n}) should equal {extra}"
+            );
+        }
+    }
+
+    // ── should_retain ────────────────────────────────────────────────────────
+
+    #[test]
+    fn test_should_retain_sink_positions() {
+        let sink = AttentionSinkEviction::new(4, 8);
+        let total = 100;
+
+        // All sink positions must be retained regardless of `total`.
+        for pos in 0..4usize {
+            assert!(
+                sink.should_retain(pos, total),
+                "sink position {pos} must be retained"
+            );
+        }
+    }
+
+    #[test]
+    fn test_should_retain_recent_positions() {
+        let sink = AttentionSinkEviction::new(4, 8);
+        let total = 100;
+
+        // Positions 92..100 (last 8) must always be retained.
+        for pos in 92..100usize {
+            assert!(
+                sink.should_retain(pos, total),
+                "recent position {pos} must be retained"
+            );
+        }
+    }
+
+    #[test]
+    fn test_should_not_retain_middle() {
+        let sink = AttentionSinkEviction::new(4, 8);
+        let total = 100;
+
+        // Positions [4, 92) are in the evictable middle band.
+        for pos in 4..92usize {
+            assert!(
+                !sink.should_retain(pos, total),
+                "middle position {pos} must not be retained when total={total}"
+            );
+        }
+    }
+
+    // ── edge cases: zero sink / zero window ──────────────────────────────────
+
+    #[test]
+    fn test_zero_sink_size() {
+        // sink_size=0 means no positions are pinned at the front.
+        // Only the last `window_size` positions are kept.
+        let sink = AttentionSinkEviction::new(0, 8);
+        let total = 20;
+
+        let evicted = sink.select_evict_positions(total);
+        // capacity=8, so 12 positions evicted: indices 0..12.
+        assert_eq!(evicted.len(), 12);
+        assert_eq!(evicted[0], 0, "first evicted position should be 0 when sink_size=0");
+
+        // The last 8 positions (12..20) must all be retained.
+        for pos in 12..20usize {
+            assert!(sink.should_retain(pos, total));
+        }
+        // Position 0 is NOT a sink and NOT in window → evicted.
+        assert!(!sink.should_retain(0, total));
+    }
+
+    #[test]
+    fn test_zero_window_size() {
+        // window_size=0 means only the sink tokens are retained.
+        let sink = AttentionSinkEviction::new(4, 0);
+        let total = 10;
+
+        let evicted = sink.select_evict_positions(total);
+        // capacity=4, so 6 positions evicted: indices 4..10.
+        assert_eq!(evicted.len(), 6);
+        assert_eq!(evicted, (4..10).collect::<Vec<usize>>());
+
+        // Sink positions still retained.
+        for pos in 0..4usize {
+            assert!(sink.should_retain(pos, total));
+        }
+        // Non-sink positions not retained.
+        for pos in 4..10usize {
+            assert!(!sink.should_retain(pos, total));
+        }
+    }
+
+    // ── large eviction ───────────────────────────────────────────────────────
+
+    #[test]
+    fn test_large_eviction() {
+        let sink = AttentionSinkEviction::new(4, 32);
+        let total = 1000;
+
+        let evicted = sink.select_evict_positions(total);
+        // capacity=36 → evict 964 positions (indices 4..968).
+        assert_eq!(evicted.len(), 964, "expected 964 evictions for 1000 positions");
+        assert_eq!(evicted[0], 4, "first evicted index should be 4 (just past sinks)");
+        assert_eq!(evicted[963], 967, "last evicted index should be 967");
+
+        // Sink positions never evicted.
+        for pos in 0..4usize {
+            assert!(!evicted.contains(&pos));
+        }
+        // Window positions never evicted.
+        for pos in 968..1000usize {
+            assert!(!evicted.contains(&pos));
+        }
+    }
+
+    #[test]
+    fn test_evicted_positions_are_sorted() {
+        let sink = AttentionSinkEviction::new(3, 5);
+
+        // Test multiple sequence lengths above capacity.
+        for total in [9, 20, 50, 100, 1000] {
+            let evicted = sink.select_evict_positions(total);
+            let mut sorted = evicted.clone();
+            sorted.sort_unstable();
+            assert_eq!(
+                evicted, sorted,
+                "evicted positions for total={total} must be sorted ascending"
+            );
+        }
+    }
+
+    // ── additional correctness checks ─────────────────────────────────────────
+
+    #[test]
+    fn test_max_capacity_arithmetic() {
+        assert_eq!(AttentionSinkEviction::new(0, 0).max_capacity(), 0);
+        assert_eq!(AttentionSinkEviction::new(1, 0).max_capacity(), 1);
+        assert_eq!(AttentionSinkEviction::new(0, 1).max_capacity(), 1);
+        assert_eq!(AttentionSinkEviction::new(4, 32).max_capacity(), 36);
+        assert_eq!(AttentionSinkEviction::new(128, 512).max_capacity(), 640);
+    }
+
+    #[test]
+    fn test_default_parameters() {
+        let sink = AttentionSinkEviction::default();
+        assert_eq!(sink.sink_size, 4);
+        assert_eq!(sink.window_size, 128);
+        assert_eq!(sink.max_capacity(), 132);
+    }
+
+    #[test]
+    fn test_select_evict_positions_matches_should_retain() {
+        let sink = AttentionSinkEviction::new(4, 8);
+        let total = 30;
+
+        let evicted = sink.select_evict_positions(total);
+
+        for pos in 0..total {
+            let retained = sink.should_retain(pos, total);
+            let in_evict_list = evicted.contains(&pos);
+            assert!(
+                retained != in_evict_list,
+                "position {pos}: should_retain={retained} but in evict list={in_evict_list} — these must disagree"
+            );
+        }
+    }
+}

crates/production/rtx-inference/src/cache/mod.rs

@@ -17,6 +17,7 @@
 //! Zobrist-hashed token prefixes to their KV page IDs, allowing requests
 //! sharing a common prefix (e.g. a system prompt) to skip recomputation.
 
+mod attention_sink;
 mod eviction;
 pub mod kv_quant;
 mod manager;
@@ -27,6 +28,7 @@ mod tracker;
 mod types;
 
 // Re-export all public types
+pub use attention_sink::AttentionSinkEviction;
 pub use eviction::AttentionScoreEviction;
 pub use kv_quant::{KvCacheQuantMode, KvQuantizer, QuantizedKvBlock};
 pub use manager::PagedKvCacheManager;

crates/production/rtx-inference/src/chunked_prefill.rs

@@ -0,0 +1,573 @@
+//! Chunked prefill scheduler — interleaves prefill chunks with decode steps.
+//!
+//! For long prompts, processing the entire prompt in one shot causes quadratic
+//! attention memory (`O(seq^2)`) and starves decode requests waiting in queue.
+//! Chunked prefill (cf. arXiv:2309.06180) solves this by:
+//!
+//! - Splitting each prompt into fixed-size *chunks* (default: 512 tokens).
+//! - Scheduling one chunk per step, optionally co-scheduling active decode
+//!   requests in the same GPU step.
+//! - Bounding first-token latency for long prompts while sustaining throughput.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_inference::chunked_prefill::{
+//!     ChunkedPrefillConfig, ChunkedPrefillScheduler, PrefillChunkState,
+//! };
+//!
+//! let config = ChunkedPrefillConfig {
+//!     chunk_size: 512,
+//!     max_decode_tokens: 128,
+//!     interleave_decode: true,
+//! };
+//!
+//! let mut scheduler = ChunkedPrefillScheduler::new(config);
+//!
+//! // Enqueue a 1024-token prompt.
+//! scheduler.enqueue(42, 1024);
+//!
+//! // First step: 512 prefill tokens + up to 128 decode tokens.
+//! let step = scheduler.next_step(8);
+//! assert_eq!(step.prefill_tokens, 512);
+//! assert_eq!(step.decode_tokens, 8); // min(128, 8)
+//! assert_eq!(step.prefill_request_id, Some(42));
+//!
+//! // Second step: remaining 512 tokens.
+//! let step = scheduler.next_step(8);
+//! assert_eq!(step.prefill_tokens, 512);
+//!
+//! // Prefill is now complete; drain and confirm.
+//! assert_eq!(scheduler.drain_completed(), 1);
+//! assert!(scheduler.is_idle());
+//! ```
+
+/// Configuration for chunked prefill scheduling.
+///
+/// All fields have sensible defaults via [`Default`]:
+/// - `chunk_size = 512`
+/// - `max_decode_tokens = 128`
+/// - `interleave_decode = true`
+#[derive(Debug, Clone)]
+pub struct ChunkedPrefillConfig {
+    /// Maximum number of tokens to process per prefill chunk.
+    ///
+    /// Smaller values reduce peak attention memory at the cost of more
+    /// scheduling overhead.  512 is a good default for 80 GB A100s.
+    pub chunk_size: usize,
+
+    /// Maximum decode tokens to process alongside a prefill chunk.
+    ///
+    /// When `interleave_decode` is `true`, each step runs at most this many
+    /// decode tokens in parallel with the prefill chunk.  The actual count is
+    /// `min(max_decode_tokens, active_decode_count)`.
+    pub max_decode_tokens: usize,
+
+    /// Whether to interleave decode requests with prefill chunks.
+    ///
+    /// Set to `false` to dedicate each GPU step entirely to prefill (useful
+    /// when decode latency is not a concern, e.g. offline batch processing).
+    pub interleave_decode: bool,
+}
+
+impl Default for ChunkedPrefillConfig {
+    fn default() -> Self {
+        Self {
+            chunk_size: 512,
+            max_decode_tokens: 128,
+            interleave_decode: true,
+        }
+    }
+}
+
+/// Progress state for a single request that is being chunked through prefill.
+///
+/// Created by [`ChunkedPrefillScheduler::enqueue`] and updated by each call
+/// to [`ChunkedPrefillScheduler::next_step`].
+#[derive(Debug, Clone)]
+pub struct PrefillChunkState {
+    /// Unique identifier supplied by the caller (mirrors request IDs in the
+    /// surrounding scheduler infrastructure).
+    pub request_id: u64,
+
+    /// Total prompt tokens for this request.
+    pub total_tokens: usize,
+
+    /// How many tokens have already been prefilled.
+    pub processed_tokens: usize,
+
+    /// `true` once `processed_tokens == total_tokens`.
+    pub is_complete: bool,
+}
+
+impl PrefillChunkState {
+    /// Create a new state for `request_id` with `total_tokens` to process.
+    ///
+    /// # Panics
+    ///
+    /// Panics in debug builds if `total_tokens == 0`, which would be a
+    /// programming error (zero-length prompts should be rejected upstream).
+    #[must_use]
+    pub fn new(request_id: u64, total_tokens: usize) -> Self {
+        debug_assert!(total_tokens > 0, "total_tokens must be > 0");
+        Self {
+            request_id,
+            total_tokens,
+            processed_tokens: 0,
+            is_complete: false,
+        }
+    }
+
+    /// Advance by `n` tokens.
+    ///
+    /// Clamps to `total_tokens` so callers do not need to guard against
+    /// overshooting.  Sets `is_complete` when all tokens are processed.
+    pub fn advance(&mut self, n: usize) {
+        self.processed_tokens = self.processed_tokens.saturating_add(n).min(self.total_tokens);
+        self.is_complete = self.processed_tokens >= self.total_tokens;
+    }
+
+    /// Number of tokens still to be prefilled.
+    #[must_use]
+    pub fn remaining(&self) -> usize {
+        self.total_tokens.saturating_sub(self.processed_tokens)
+    }
+
+    /// Fraction of the prompt that has been processed, in `0.0..=1.0`.
+    ///
+    /// Returns `1.0` for zero-length prompts to avoid a division-by-zero.
+    #[must_use]
+    pub fn progress(&self) -> f32 {
+        if self.total_tokens == 0 {
+            return 1.0;
+        }
+        self.processed_tokens as f32 / self.total_tokens as f32
+    }
+
+    /// Size of the next chunk to schedule, given `config`.
+    ///
+    /// Always `<= config.chunk_size` and `<= self.remaining()`.
+    /// Returns `0` when `is_complete`.
+    #[must_use]
+    pub fn next_chunk_size(&self, config: &ChunkedPrefillConfig) -> usize {
+        self.remaining().min(config.chunk_size)
+    }
+}
+
+/// One scheduling decision returned by [`ChunkedPrefillScheduler::next_step`].
+///
+/// The executing engine should:
+/// 1. Run `prefill_tokens` tokens of prompt `prefill_request_id`.
+/// 2. Run `decode_tokens` tokens for the currently active decode requests.
+#[derive(Debug, Clone)]
+pub struct ChunkedStep {
+    /// Request whose prefill chunk is scheduled this step.
+    ///
+    /// `None` when there is no active prefill (pure decode step).
+    pub prefill_request_id: Option<u64>,
+
+    /// Number of prefill tokens to process this step.
+    ///
+    /// `0` when there is no active prefill request.
+    pub prefill_tokens: usize,
+
+    /// Number of decode tokens to process this step.
+    ///
+    /// `0` when `interleave_decode` is `false` or `active_decode_count == 0`.
+    pub decode_tokens: usize,
+
+    /// Monotonically increasing step counter (zero-based).
+    pub step_index: usize,
+}
+
+/// Scheduler that splits long prompts into fixed-size chunks and optionally
+/// interleaves decode requests within each step.
+///
+/// # Design notes
+///
+/// - Uses a `Vec<PrefillChunkState>` (not a `VecDeque`) because the active
+///   set is typically very small (< 10 requests).  FIFO ordering is preserved
+///   by always picking `active[0]` as the head.
+/// - Completed states are retained until [`drain_completed`] is called, which
+///   mirrors the explicit lifecycle in production schedulers.
+/// - All methods are `&mut self`; the struct is **not** `Send` by default
+///   because its callers embed it inside a larger struct that owns the lock.
+pub struct ChunkedPrefillScheduler {
+    config: ChunkedPrefillConfig,
+    /// Requests currently being chunked (includes completed ones until
+    /// [`drain_completed`] is called).
+    active: Vec<PrefillChunkState>,
+    /// Monotonically increasing step counter.
+    step_counter: usize,
+    /// Cumulative number of prefill chunks emitted.
+    chunks_scheduled: usize,
+    /// Cumulative number of tokens prefilled across all requests.
+    tokens_prefilled: usize,
+}
+
+impl ChunkedPrefillScheduler {
+    /// Create a new scheduler with the given configuration.
+    #[must_use]
+    pub fn new(config: ChunkedPrefillConfig) -> Self {
+        Self {
+            config,
+            active: Vec::new(),
+            step_counter: 0,
+            chunks_scheduled: 0,
+            tokens_prefilled: 0,
+        }
+    }
+
+    /// Enqueue `request_id` for chunked prefill of `total_tokens` prompt tokens.
+    ///
+    /// The request will be scheduled FIFO after any currently active prefill
+    /// requests.
+    pub fn enqueue(&mut self, request_id: u64, total_tokens: usize) {
+        self.active.push(PrefillChunkState::new(request_id, total_tokens));
+    }
+
+    /// Return the next [`ChunkedStep`] to execute.
+    ///
+    /// **Prefill logic**: picks the first incomplete request in `active` and
+    /// schedules its next chunk.  The state is advanced immediately so the
+    /// caller does not need to report completion back.
+    ///
+    /// **Decode interleave**: when `config.interleave_decode` is `true` and
+    /// `active_decode_count > 0`, `decode_tokens` is set to
+    /// `min(config.max_decode_tokens, active_decode_count)`.
+    ///
+    /// Returns a **pure decode step** (prefill fields zeroed, `prefill_request_id
+    /// = None`) when there is no active prefill work.
+    pub fn next_step(&mut self, active_decode_count: usize) -> ChunkedStep {
+        let step_index = self.step_counter;
+        self.step_counter += 1;
+
+        // Decode token count is independent of whether prefill is active.
+        let decode_tokens = if self.config.interleave_decode {
+            active_decode_count.min(self.config.max_decode_tokens)
+        } else {
+            0
+        };
+
+        // Find the first incomplete prefill request.
+        let head = self.active.iter_mut().find(|s| !s.is_complete);
+
+        match head {
+            Some(state) => {
+                let chunk = state.next_chunk_size(&self.config);
+                let request_id = state.request_id;
+                state.advance(chunk);
+
+                self.chunks_scheduled += 1;
+                self.tokens_prefilled += chunk;
+
+                ChunkedStep {
+                    prefill_request_id: Some(request_id),
+                    prefill_tokens: chunk,
+                    decode_tokens,
+                    step_index,
+                }
+            }
+            None => {
+                // No active prefill — pure decode step.
+                ChunkedStep {
+                    prefill_request_id: None,
+                    prefill_tokens: 0,
+                    decode_tokens,
+                    step_index,
+                }
+            }
+        }
+    }
+
+    /// Number of requests currently tracked (including completed ones not yet
+    /// drained).
+    #[must_use]
+    pub fn active_count(&self) -> usize {
+        self.active.len()
+    }
+
+    /// `true` when there are no incomplete prefill requests.
+    ///
+    /// Note: returns `true` on an empty scheduler (nothing to do).
+    #[must_use]
+    pub fn is_idle(&self) -> bool {
+        self.active.iter().all(|s| s.is_complete)
+    }
+
+    /// Remove all completed requests and return how many were removed.
+    pub fn drain_completed(&mut self) -> usize {
+        let before = self.active.len();
+        self.active.retain(|s| !s.is_complete);
+        before - self.active.len()
+    }
+
+    /// Cumulative number of prefill chunks emitted by [`next_step`].
+    #[must_use]
+    pub fn chunks_scheduled(&self) -> usize {
+        self.chunks_scheduled
+    }
+
+    /// Cumulative number of tokens prefilled across all requests.
+    #[must_use]
+    pub fn tokens_prefilled(&self) -> usize {
+        self.tokens_prefilled
+    }
+}
+
+// ── Unit tests ────────────────────────────────────────────────────────────────
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn default_config() -> ChunkedPrefillConfig {
+        ChunkedPrefillConfig::default()
+    }
+
+    // ------------------------------------------------------------------
+    // PrefillChunkState tests
+    // ------------------------------------------------------------------
+
+    /// A 100-token request with chunk_size=512 should complete in exactly one
+    /// step.
+    #[test]
+    fn test_single_short_request_one_chunk() {
+        let config = default_config(); // chunk_size = 512
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        scheduler.enqueue(1, 100);
+
+        let step = scheduler.next_step(0);
+        assert_eq!(step.prefill_request_id, Some(1));
+        assert_eq!(step.prefill_tokens, 100);
+        assert_eq!(step.step_index, 0);
+
+        // After one step the request is complete.
+        assert!(scheduler.is_idle());
+        assert_eq!(scheduler.tokens_prefilled(), 100);
+        assert_eq!(scheduler.chunks_scheduled(), 1);
+    }
+
+    /// A 1500-token request with chunk_size=512 must require exactly 3 steps:
+    /// 512 + 512 + 476 = 1500.
+    #[test]
+    fn test_long_request_multiple_chunks() {
+        let config = default_config(); // chunk_size = 512
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        scheduler.enqueue(7, 1500);
+
+        let s0 = scheduler.next_step(0);
+        assert_eq!(s0.prefill_tokens, 512);
+        assert!(!scheduler.is_idle());
+
+        let s1 = scheduler.next_step(0);
+        assert_eq!(s1.prefill_tokens, 512);
+        assert!(!scheduler.is_idle());
+
+        let s2 = scheduler.next_step(0);
+        assert_eq!(s2.prefill_tokens, 476); // 1500 - 512 - 512
+        assert!(scheduler.is_idle());
+
+        assert_eq!(scheduler.tokens_prefilled(), 1500);
+        assert_eq!(scheduler.chunks_scheduled(), 3);
+    }
+
+    /// After advancing 256 of 1024 tokens, progress() should be exactly 0.25.
+    #[test]
+    fn test_progress_fraction() {
+        let mut state = PrefillChunkState::new(99, 1024);
+        state.advance(256);
+        let got = state.progress();
+        assert!(
+            (got - 0.25_f32).abs() < 1e-6,
+            "expected 0.25, got {got}"
+        );
+    }
+
+    /// remaining() must equal total_tokens - processed_tokens after advancing.
+    #[test]
+    fn test_remaining_tokens() {
+        let mut state = PrefillChunkState::new(3, 800);
+        state.advance(300);
+        assert_eq!(state.remaining(), 500);
+        state.advance(500);
+        assert_eq!(state.remaining(), 0);
+        assert!(state.is_complete);
+    }
+
+    /// When only a partial chunk remains, next_chunk_size must return the
+    /// smaller remainder rather than the full chunk_size.
+    #[test]
+    fn test_next_chunk_size_last_chunk() {
+        let config = ChunkedPrefillConfig {
+            chunk_size: 512,
+            ..Default::default()
+        };
+        let mut state = PrefillChunkState::new(5, 700);
+        state.advance(512); // first chunk consumed
+        let next = state.next_chunk_size(&config);
+        assert_eq!(next, 188); // 700 - 512 = 188 < 512
+    }
+
+    // ------------------------------------------------------------------
+    // Decode interleave tests
+    // ------------------------------------------------------------------
+
+    /// With interleave_decode=true, decode_tokens should be
+    /// min(max_decode_tokens, active_decode_count).
+    #[test]
+    fn test_interleave_decode_tokens() {
+        let config = ChunkedPrefillConfig {
+            chunk_size: 512,
+            max_decode_tokens: 128,
+            interleave_decode: true,
+        };
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        scheduler.enqueue(10, 1000);
+
+        let step = scheduler.next_step(5); // 5 active decode requests
+        assert_eq!(step.decode_tokens, 5); // min(128, 5)
+
+        let step2 = scheduler.next_step(200); // more than max_decode_tokens
+        assert_eq!(step2.decode_tokens, 128); // clamped to max
+    }
+
+    /// With interleave_decode=false, decode_tokens must always be 0.
+    #[test]
+    fn test_no_decode_when_disabled() {
+        let config = ChunkedPrefillConfig {
+            chunk_size: 512,
+            max_decode_tokens: 128,
+            interleave_decode: false,
+        };
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        scheduler.enqueue(20, 600);
+
+        for _ in 0..2 {
+            let step = scheduler.next_step(99);
+            assert_eq!(step.decode_tokens, 0, "decode must be 0 when disabled");
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // Lifecycle / bookkeeping tests
+    // ------------------------------------------------------------------
+
+    /// After a request is fully prefilled, drain_completed() should remove it
+    /// and return 1.
+    #[test]
+    fn test_drain_completed_removes_done_requests() {
+        let config = default_config();
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        scheduler.enqueue(42, 100); // completes in one step (< 512)
+
+        scheduler.next_step(0); // completes the request
+        assert!(scheduler.is_idle());
+
+        let removed = scheduler.drain_completed();
+        assert_eq!(removed, 1);
+        assert_eq!(scheduler.active_count(), 0);
+    }
+
+    /// active_count() must reflect the number of tracked requests (including
+    /// completed ones not yet drained).
+    #[test]
+    fn test_active_count() {
+        let config = default_config();
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        assert_eq!(scheduler.active_count(), 0);
+
+        scheduler.enqueue(1, 50);
+        scheduler.enqueue(2, 50);
+        assert_eq!(scheduler.active_count(), 2);
+
+        // Complete both requests.
+        scheduler.next_step(0); // req 1 done
+        scheduler.next_step(0); // req 2 done
+
+        // Still 2 until explicitly drained.
+        assert_eq!(scheduler.active_count(), 2);
+        scheduler.drain_completed();
+        assert_eq!(scheduler.active_count(), 0);
+    }
+
+    /// is_idle() must return true once all requests have been drained.
+    #[test]
+    fn test_idle_after_all_complete() {
+        let config = default_config();
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        scheduler.enqueue(55, 256);
+
+        assert!(!scheduler.is_idle());
+        scheduler.next_step(0); // chunk_size=512 > 256, so completes in one step
+        assert!(scheduler.is_idle());
+
+        scheduler.drain_completed();
+        // Empty scheduler is also idle.
+        assert!(scheduler.is_idle());
+    }
+
+    /// step_index must increment by 1 for every call to next_step.
+    #[test]
+    fn test_step_index_increments() {
+        let config = default_config();
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        scheduler.enqueue(100, 2000);
+
+        for expected_idx in 0..5_usize {
+            let step = scheduler.next_step(0);
+            assert_eq!(
+                step.step_index, expected_idx,
+                "step_index should be {expected_idx}"
+            );
+        }
+    }
+
+    /// tokens_prefilled() must equal the sum of all chunk sizes emitted.
+    #[test]
+    fn test_cumulative_tokens_prefilled() {
+        let config = ChunkedPrefillConfig {
+            chunk_size: 200,
+            ..Default::default()
+        };
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+        // Request A: 400 tokens → 2 chunks of 200 each.
+        scheduler.enqueue(1, 400);
+        // Request B: 300 tokens → 2 chunks of 200+100.
+        scheduler.enqueue(2, 300);
+
+        // Drive all chunks to completion.
+        while !scheduler.is_idle() {
+            scheduler.next_step(0);
+        }
+
+        assert_eq!(scheduler.tokens_prefilled(), 700); // 400 + 300
+        assert_eq!(scheduler.chunks_scheduled(), 4);   // 2 + 2
+    }
+
+    // ------------------------------------------------------------------
+    // Extra edge-case tests
+    // ------------------------------------------------------------------
+
+    /// advance() must clamp and not overflow even if n >> total_tokens.
+    #[test]
+    fn test_advance_clamps_to_total() {
+        let mut state = PrefillChunkState::new(9, 100);
+        state.advance(999); // deliberately overshooting
+        assert_eq!(state.processed_tokens, 100);
+        assert!(state.is_complete);
+        assert_eq!(state.remaining(), 0);
+    }
+
+    /// A pure decode step (no enqueued prefill) must have zero prefill fields.
+    #[test]
+    fn test_pure_decode_step_when_no_prefill() {
+        let config = default_config();
+        let mut scheduler = ChunkedPrefillScheduler::new(config);
+
+        let step = scheduler.next_step(10);
+        assert_eq!(step.prefill_request_id, None);
+        assert_eq!(step.prefill_tokens, 0);
+        assert_eq!(step.decode_tokens, 10);
+    }
+}

crates/production/rtx-inference/src/lib.rs

@@ -17,6 +17,8 @@
 
 pub mod batch_processor;
 pub mod cache;
+pub mod chunked_prefill;
+pub use chunked_prefill::{ChunkedPrefillConfig, ChunkedPrefillScheduler, ChunkedStep, PrefillChunkState};
 pub mod inference_graph;
 pub use inference_graph::{InferenceGraphCapture, StepMode};
 pub mod engine;
@@ -30,8 +32,8 @@ pub mod speculative;
 
 // Re-export key types for convenience
 pub use cache::{
-    AttentionScoreEviction, CacheKey, CachePage, CacheStats, EvictionPolicy, KvCacheConfig,
+    AttentionScoreEviction, AttentionSinkEviction, CacheKey, CachePage, CacheStats, EvictionPolicy,
-    MemoryTier, PageId, PagedKvCache, PagedKvCacheManager, PrefixIndex,
+    KvCacheConfig, MemoryTier, PageId, PagedKvCache, PagedKvCacheManager, PrefixIndex,
 };
 pub use engine::{
     HealthStatus, InferenceEngine, InferenceEngineConfig, MemoryStats, ModelConfig, ModelHealth,

crates/training/rtx-transformers/src/optimizers/layer_lr_decay.rs

@@ -0,0 +1,563 @@
+//! Per-layer learning rate decay for discriminative fine-tuning.
+//!
+//! Implements the discriminative fine-tuning technique from Howard & Ruder (2018)
+//! "Universal Language Model Fine-tuning for Text Classification" (ULMFiT).
+//!
+//! Lower transformer layers (closer to input embeddings) receive smaller learning
+//! rates, while deeper layers receive the full base learning rate.  This stabilises
+//! fine-tuning because lower layers encode general features that should change
+//! slowly, whereas higher layers are task-specific and can tolerate larger updates.
+//!
+//! # Formula
+//!
+//! ```text
+//! lr(depth) = clamp(base_lr * decay_rate ^ (num_layers - 1 - depth), min_lr, base_lr)
+//! ```
+//!
+//! Where `depth = 0` is the shallowest layer (closest to input/embeddings) and
+//! `depth = num_layers - 1` is the deepest layer (closest to the output head).
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::optimizers::layer_lr_decay::{LayerLrDecayConfig, LayerLrDecayBuilder};
+//!
+//! let config = LayerLrDecayConfig::new(1e-3, 0.9, 12);
+//! let builder = LayerLrDecayBuilder::new(config);
+//!
+//! let param_names = &[
+//!     "embeddings.weight",
+//!     "layer.0.weight",
+//!     "layer.11.weight",
+//!     "lm_head.weight",
+//! ];
+//! let groups = builder.assign(param_names);
+//! for g in &groups {
+//!     println!("{}: lr = {:.2e}", g.name, g.lr);
+//! }
+//! ```
+
+// ─────────────────────────────────────────────────────────────────────────────
+// LayerLrDecayConfig
+// ─────────────────────────────────────────────────────────────────────────────
+
+/// Configuration for per-layer learning rate decay.
+///
+/// Controls how the learning rate decays from the output layer toward the
+/// embedding/input layer of a transformer model.
+#[derive(Debug, Clone)]
+pub struct LayerLrDecayConfig {
+    /// Base (deepest/top layer) learning rate.
+    pub base_lr: f32,
+    /// Decay multiplier applied per step toward the input (`0.0..=1.0`).
+    ///
+    /// `lr_for_layer(depth) = base_lr * decay_rate ^ (num_layers - 1 - depth)`
+    pub decay_rate: f32,
+    /// Total number of transformer layers (excluding embedding/head).
+    pub num_layers: usize,
+    /// Minimum LR floor — computed values are clamped to at least this value.
+    ///
+    /// Defaults to `1e-7`.
+    pub min_lr: f32,
+}
+
+impl LayerLrDecayConfig {
+    /// Create a new configuration with the given base LR, decay rate, and layer count.
+    ///
+    /// `min_lr` is initialised to `1e-7`.
+    #[must_use]
+    pub fn new(base_lr: f32, decay_rate: f32, num_layers: usize) -> Self {
+        Self {
+            base_lr,
+            decay_rate,
+            num_layers,
+            min_lr: 1e-7,
+        }
+    }
+
+    /// Compute the learning rate for a given layer `depth`.
+    ///
+    /// `depth = 0` is the shallowest layer (closest to input/embeddings).
+    /// `depth = num_layers - 1` is the deepest layer (closest to the output head).
+    ///
+    /// The result is clamped to `[min_lr, base_lr]`.
+    #[must_use]
+    pub fn lr_for_layer(&self, depth: usize) -> f32 {
+        if self.num_layers == 0 {
+            return self.base_lr;
+        }
+        // exponent = 0 at the top layer (depth = num_layers-1), increases toward input
+        let exponent = (self.num_layers - 1).saturating_sub(depth) as i32;
+        let raw = self.base_lr * self.decay_rate.powi(exponent);
+        raw.max(self.min_lr).min(self.base_lr)
+    }
+
+    /// Compute the learning rate for the embedding layer.
+    ///
+    /// Embeddings are treated as one level below `depth = 0`, so the exponent is
+    /// `num_layers` (one more than for `depth = 0`).
+    #[must_use]
+    pub fn lr_for_embeddings(&self) -> f32 {
+        let exponent = self.num_layers as i32;
+        let raw = self.base_lr * self.decay_rate.powi(exponent);
+        raw.max(self.min_lr).min(self.base_lr)
+    }
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// ParamGroupLR
+// ─────────────────────────────────────────────────────────────────────────────
+
+/// Assignment of a learning rate to a named parameter group.
+#[derive(Debug, Clone, PartialEq)]
+pub struct ParamGroupLR {
+    /// Parameter group name (e.g. `"layer.0"`, `"layer.11"`, `"embeddings"`).
+    pub name: String,
+    /// Assigned learning rate.
+    pub lr: f32,
+    /// Layer depth this group corresponds to (`None` for embeddings/head).
+    pub depth: Option<usize>,
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Layer name parsing
+// ─────────────────────────────────────────────────────────────────────────────
+
+/// Layer prefixes recognised when extracting a depth index from a parameter name.
+const LAYER_PREFIXES: &[&str] = &["layer", "layers", "h", "blocks", "block", "encoder_layer"];
+
+/// Try to extract a numeric layer depth from a dot-separated parameter name.
+///
+/// Strategy: scan each consecutive pair of dot-separated segments where the first
+/// segment matches one of [`LAYER_PREFIXES`] and the second parses as a `usize`.
+/// Bracket notation (`name[N]`) is normalised to dots before scanning.
+///
+/// Returns `Some(N)` on the first match, `None` if no layer depth is found.
+fn extract_layer_depth(name: &str) -> Option<usize> {
+    // Normalise bracket notation: "h[3].weight" → "h.3..weight"
+    let normalised: String = name
+        .chars()
+        .map(|c| match c {
+            '[' | ']' => '.',
+            other => other,
+        })
+        .collect();
+
+    let parts: Vec<&str> = normalised
+        .split('.')
+        .filter(|s| !s.is_empty())
+        .collect();
+
+    for window in parts.windows(2) {
+        let prefix = window[0];
+        let candidate = window[1];
+        if LAYER_PREFIXES.contains(&prefix) {
+            if let Ok(n) = candidate.parse::<usize>() {
+                return Some(n);
+            }
+        }
+    }
+
+    None
+}
+
+/// Returns `true` if the parameter name refers to an embedding layer.
+fn is_embedding_name(name: &str) -> bool {
+    let lower = name.to_lowercase();
+    lower.contains("embed")
+        || lower.contains("wte")
+        || lower.contains("wpe")
+        || lower.contains("token_embedding")
+        || lower.contains("position_embedding")
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// LayerLrDecayBuilder
+// ─────────────────────────────────────────────────────────────────────────────
+
+/// Builds per-layer LR assignments from a list of parameter names.
+///
+/// Matches names against common transformer layer patterns to extract a numeric
+/// depth, then applies [`LayerLrDecayConfig::lr_for_layer`].  Embedding-like
+/// names use [`LayerLrDecayConfig::lr_for_embeddings`]; everything else
+/// (heads, final norms, classifiers) receives `base_lr`.
+pub struct LayerLrDecayBuilder {
+    config: LayerLrDecayConfig,
+}
+
+impl LayerLrDecayBuilder {
+    /// Create a new builder from the given config.
+    #[must_use]
+    pub fn new(config: LayerLrDecayConfig) -> Self {
+        Self { config }
+    }
+
+    /// Assign learning rates to a slice of parameter names.
+    ///
+    /// Returns one [`ParamGroupLR`] per input name, in the same order.
+    #[must_use]
+    pub fn assign(&self, param_names: &[&str]) -> Vec<ParamGroupLR> {
+        param_names
+            .iter()
+            .map(|&name| self.assign_one(name))
+            .collect()
+    }
+
+    /// Deduplicated unique LR assignments ordered by first occurrence.
+    ///
+    /// Two assignments are considered the same when they share the same `depth`
+    /// value (both `Some(N)` with equal `N`, or both `None`) **and** the same
+    /// LR bits.  At most `num_layers + 2` entries are returned (one per
+    /// numbered layer, one for embeddings, one for the head).
+    #[must_use]
+    pub fn unique_groups(&self, param_names: &[&str]) -> Vec<ParamGroupLR> {
+        let all = self.assign(param_names);
+
+        // Use a deduplication key of (depth, lr_bits).  Small N → Vec scan is fine.
+        let mut seen: Vec<(Option<usize>, u32)> = Vec::new();
+        let mut result: Vec<ParamGroupLR> = Vec::new();
+
+        for group in all {
+            let key = (group.depth, group.lr.to_bits());
+            if !seen.contains(&key) {
+                seen.push(key);
+                result.push(group);
+            }
+        }
+
+        result
+    }
+
+    /// Assign a single parameter name to its learning rate group.
+    fn assign_one(&self, name: &str) -> ParamGroupLR {
+        if let Some(depth) = extract_layer_depth(name) {
+            // Clamp depth to valid range in case the name exceeds configured layers.
+            let effective_depth = if self.config.num_layers == 0 {
+                0
+            } else {
+                depth.min(self.config.num_layers - 1)
+            };
+            ParamGroupLR {
+                name: name.to_string(),
+                lr: self.config.lr_for_layer(effective_depth),
+                depth: Some(effective_depth),
+            }
+        } else if is_embedding_name(name) {
+            ParamGroupLR {
+                name: name.to_string(),
+                lr: self.config.lr_for_embeddings(),
+                depth: None,
+            }
+        } else {
+            // Head, final norm, bias terms not inside a named layer → full base LR.
+            ParamGroupLR {
+                name: name.to_string(),
+                lr: self.config.base_lr,
+                depth: None,
+            }
+        }
+    }
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// LayerLrScheduler
+// ─────────────────────────────────────────────────────────────────────────────
+
+/// Wraps [`LayerLrDecayBuilder`] with an outer scalar multiplier.
+///
+/// The multiplier is typically updated each training step from an outer cosine
+/// or linear decay schedule.  All effective LRs equal the config-derived LR
+/// multiplied by the current multiplier.
+pub struct LayerLrScheduler {
+    builder: LayerLrDecayBuilder,
+    multiplier: f32,
+}
+
+impl LayerLrScheduler {
+    /// Create a new scheduler with `multiplier = 1.0`.
+    #[must_use]
+    pub fn new(config: LayerLrDecayConfig) -> Self {
+        Self {
+            builder: LayerLrDecayBuilder::new(config),
+            multiplier: 1.0,
+        }
+    }
+
+    /// Update the global multiplier (call from the outer LR scheduler each step).
+    pub fn set_multiplier(&mut self, m: f32) {
+        self.multiplier = m;
+    }
+
+    /// Effective LR for a given layer `depth` after applying the multiplier.
+    #[must_use]
+    pub fn effective_lr(&self, depth: usize) -> f32 {
+        self.builder.config.lr_for_layer(depth) * self.multiplier
+    }
+
+    /// Effective LR for the embedding layer after applying the multiplier.
+    #[must_use]
+    pub fn effective_lr_embeddings(&self) -> f32 {
+        self.builder.config.lr_for_embeddings() * self.multiplier
+    }
+
+    /// All effective LRs from `depth = 0` (shallowest) to `depth = num_layers - 1`.
+    #[must_use]
+    pub fn all_lrs(&self) -> Vec<f32> {
+        (0..self.builder.config.num_layers)
+            .map(|d| self.effective_lr(d))
+            .collect()
+    }
+}
+
+// ─────────────────────────────────────────────────────────────────────────────
+// Tests
+// ─────────────────────────────────────────────────────────────────────────────
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const BASE_LR: f32 = 1e-3;
+    const DECAY: f32 = 0.9;
+    const N: usize = 12;
+
+    fn cfg() -> LayerLrDecayConfig {
+        LayerLrDecayConfig::new(BASE_LR, DECAY, N)
+    }
+
+    // 1. Top layer (deepest) gets exactly base_lr ────────────────────────────
+
+    #[test]
+    fn test_lr_for_top_layer() {
+        let config = cfg();
+        let lr = config.lr_for_layer(N - 1);
+        // exponent = (N-1) - (N-1) = 0 → decay^0 = 1 → base_lr
+        assert!(
+            (lr - BASE_LR).abs() < 1e-9,
+            "top layer lr should equal base_lr={BASE_LR}, got {lr}"
+        );
+    }
+
+    // 2. Bottom layer (shallowest) gets base_lr * decay^(N-1) ────────────────
+
+    #[test]
+    fn test_lr_for_bottom_layer() {
+        let config = cfg();
+        let lr = config.lr_for_layer(0);
+        let expected = BASE_LR * DECAY.powi((N - 1) as i32);
+        assert!(
+            (lr - expected).abs() < 1e-9,
+            "bottom layer lr should be {expected}, got {lr}"
+        );
+    }
+
+    // 3. LR strictly increases (or stays equal) with depth ───────────────────
+
+    #[test]
+    fn test_lr_decay_monotone_increasing_with_depth() {
+        let config = LayerLrDecayConfig::new(1e-3, 0.85, 8);
+        for d in 0..7_usize {
+            let lr_shallow = config.lr_for_layer(d);
+            let lr_deep = config.lr_for_layer(d + 1);
+            assert!(
+                lr_deep >= lr_shallow,
+                "lr should be non-decreasing with depth: lr({d})={lr_shallow} > lr({})={lr_deep}",
+                d + 1
+            );
+        }
+    }
+
+    // 4. Very small decay + many layers → clamped to min_lr ─────────────────
+
+    #[test]
+    fn test_lr_clamped_to_min() {
+        let mut config = LayerLrDecayConfig::new(1e-3, 0.1, 50);
+        config.min_lr = 1e-7;
+        // depth=0 → 1e-3 * 0.1^49 ≈ 1e-52 → clamped to 1e-7
+        let lr = config.lr_for_layer(0);
+        assert!(
+            (lr - 1e-7_f32).abs() < 1e-12,
+            "should be clamped to min_lr=1e-7, got {lr}"
+        );
+    }
+
+    // 5. Embeddings: one step below layer 0 → layer0_lr * decay_rate ─────────
+
+    #[test]
+    fn test_lr_for_embeddings() {
+        let config = cfg();
+        let embed_lr = config.lr_for_embeddings();
+        let layer0_lr = config.lr_for_layer(0);
+        // embed exponent = N, layer0 exponent = N-1 → embed = layer0 * decay
+        let expected = layer0_lr * DECAY;
+        assert!(
+            (embed_lr - expected).abs() < 1e-9,
+            "embed lr should be {expected}, got {embed_lr}"
+        );
+        assert!(
+            embed_lr <= layer0_lr,
+            "embed lr {embed_lr} should be <= layer0 lr {layer0_lr}"
+        );
+    }
+
+    // 6. decay_rate = 1.0 → every layer gets base_lr ─────────────────────────
+
+    #[test]
+    fn test_decay_rate_one_all_same() {
+        let config = LayerLrDecayConfig::new(5e-4, 1.0, 24);
+        for d in 0..24_usize {
+            let lr = config.lr_for_layer(d);
+            assert!(
+                (lr - 5e-4_f32).abs() < 1e-9,
+                "with decay=1.0 layer {d} should get base_lr=5e-4, got {lr}"
+            );
+        }
+    }
+
+    // 7. Named param "transformer.layer.3.weight" → depth-3 LR ───────────────
+
+    #[test]
+    fn test_assign_layer_names() {
+        let config = LayerLrDecayConfig::new(1e-3, 0.9, 12);
+        let builder = LayerLrDecayBuilder::new(config.clone());
+        let groups = builder.assign(&["transformer.layer.3.weight"]);
+        assert_eq!(groups.len(), 1);
+        let g = &groups[0];
+        assert_eq!(g.depth, Some(3), "depth should be 3");
+        let expected = config.lr_for_layer(3);
+        assert!(
+            (g.lr - expected).abs() < 1e-9,
+            "layer.3 lr should be {expected}, got {}",
+            g.lr
+        );
+    }
+
+    // 8. Embedding parameter → embedding LR ───────────────────────────────────
+
+    #[test]
+    fn test_assign_embed_name() {
+        let config = cfg();
+        let expected_embed_lr = config.lr_for_embeddings();
+        let builder = LayerLrDecayBuilder::new(config);
+        let groups = builder.assign(&["embeddings.weight"]);
+        assert_eq!(groups.len(), 1);
+        let g = &groups[0];
+        assert!(g.depth.is_none(), "embed should have no depth, got {:?}", g.depth);
+        assert!(
+            (g.lr - expected_embed_lr).abs() < 1e-9,
+            "embed lr should be {expected_embed_lr}, got {}",
+            g.lr
+        );
+    }
+
+    // 9. Head/output parameter → base_lr ──────────────────────────────────────
+
+    #[test]
+    fn test_assign_head_name() {
+        let config = cfg();
+        let builder = LayerLrDecayBuilder::new(config);
+        let groups = builder.assign(&["lm_head.weight"]);
+        assert_eq!(groups.len(), 1);
+        let g = &groups[0];
+        assert!(g.depth.is_none(), "head should have no depth");
+        assert!(
+            (g.lr - BASE_LR).abs() < 1e-9,
+            "head lr should be base_lr={BASE_LR}, got {}",
+            g.lr
+        );
+    }
+
+    // 10. unique_groups yields at most num_layers + 2 entries ─────────────────
+
+    #[test]
+    fn test_unique_groups_count() {
+        let config = LayerLrDecayConfig::new(1e-3, 0.9, 4);
+        let builder = LayerLrDecayBuilder::new(config);
+        let names: &[&str] = &[
+            "embeddings.weight",
+            "layer.0.weight",
+            "layer.0.bias",   // duplicate depth 0 → should NOT add a new group
+            "layer.1.weight",
+            "layer.2.weight",
+            "layer.3.weight",
+            "lm_head.weight",
+            "norm.weight",    // head-level (no layer depth, no embed keyword)
+        ];
+        let groups = builder.unique_groups(names);
+        assert!(
+            groups.len() <= 6,
+            "unique groups should be <= num_layers+2=6, got {}",
+            groups.len()
+        );
+    }
+
+    // 11. Multiplier = 0.5 scales every layer's effective LR by 0.5 ───────────
+
+    #[test]
+    fn test_scheduler_multiplier_scales_all() {
+        let config = LayerLrDecayConfig::new(1e-3, 0.9, 12);
+        let mut scheduler = LayerLrScheduler::new(config.clone());
+        scheduler.set_multiplier(0.5);
+        for d in 0..N {
+            let effective = scheduler.effective_lr(d);
+            let base = config.lr_for_layer(d);
+            let expected = base * 0.5;
+            assert!(
+                (effective - expected).abs() < 1e-10,
+                "depth {d}: effective={effective} should equal base*0.5={expected}"
+            );
+        }
+    }
+
+    // 12. Scheduler ordering: deeper layer → higher effective LR ──────────────
+
+    #[test]
+    fn test_scheduler_effective_lr_ordering() {
+        let config = LayerLrDecayConfig::new(1e-3, 0.85, 8);
+        let mut scheduler = LayerLrScheduler::new(config);
+        scheduler.set_multiplier(0.7);
+        let lrs = scheduler.all_lrs();
+        assert_eq!(lrs.len(), 8);
+        for i in 0..7_usize {
+            assert!(
+                lrs[i + 1] >= lrs[i],
+                "effective lr must be non-decreasing: lrs[{i}]={} > lrs[{}]={}",
+                lrs[i],
+                i + 1,
+                lrs[i + 1]
+            );
+        }
+    }
+
+    // Bonus A: Bracket notation "transformer.h[5].mlp.weight" → depth 5 ───────
+
+    #[test]
+    fn test_bracket_notation_parsed() {
+        let config = cfg();
+        let builder = LayerLrDecayBuilder::new(config.clone());
+        let groups = builder.assign(&["transformer.h[5].mlp.weight"]);
+        assert_eq!(groups[0].depth, Some(5));
+        let expected = config.lr_for_layer(5);
+        assert!(
+            (groups[0].lr - expected).abs() < 1e-9,
+            "bracket depth 5 should give lr {expected}, got {}",
+            groups[0].lr
+        );
+    }
+
+    // Bonus B: "blocks.7.attn.weight" (GPT-NeoX / Mamba style) → depth 7 ─────
+
+    #[test]
+    fn test_blocks_prefix_parsed() {
+        let config = cfg();
+        let builder = LayerLrDecayBuilder::new(config.clone());
+        let groups = builder.assign(&["blocks.7.attn.weight"]);
+        assert_eq!(groups[0].depth, Some(7));
+        let expected = config.lr_for_layer(7);
+        assert!(
+            (groups[0].lr - expected).abs() < 1e-9,
+            "blocks depth 7 should give lr {expected}, got {}",
+            groups[0].lr
+        );
+    }
+}

crates/training/rtx-transformers/src/optimizers/mod.rs

@@ -14,6 +14,7 @@ use tracing::debug;
 
 pub mod adam;
 pub mod adamw;
+pub mod layer_lr_decay;
 
 #[cfg(all(test, feature = "disabled_tests"))]
 pub mod adam_test;
@@ -60,6 +61,7 @@ pub use crate::tensor_bridge::{TensorBridge, TensorBridgeStatic, TensorCompat};
 pub use adam::AdamOptimizer;
 pub use adamw::AdamWOptimizer;
 pub use galore::{GaLoreAdamW, GaLoreConfig, GaLoreParamState};
+pub use layer_lr_decay::{LayerLrDecayBuilder, LayerLrDecayConfig, LayerLrScheduler, ParamGroupLR};
 
 // ============================================================================
 // PHASE 2: ADVANCED OPTIMIZER EXPORTS (Temporarily disabled)