feat(batch11): WSD LR scheduler, KV CPU offloading, GQA KV head expansion

- WSD scheduler (Warmup-Stable-Decay / trapezoidal): linear warmup → constant
  plateau → cosine/linear/sqrt decay; extend_stable() adds steps mid-run without
  restart; phase_at()/decay_progress() introspection; 26 tests + 2 doctests
- KV CPU offloading: KvCpuOffloadManager LRU-based GPU→CPU page spill with
  on-demand prefetch; insert() auto-offloads when at gpu_page_limit; stats()
  with hit rate and utilization; 14 tests
- GQA KV head expansion: GqaConfig validates num_q_heads/num_kv_heads divisibility;
  expand_kv_heads() tiles KV [batch,kv_heads,seq,dim]→[batch,q_heads,seq,dim];
  gqa_attention_cpu() with numerically stable softmax + causal mask; 15 tests

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

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

@@ -0,0 +1,631 @@
+//! KV cache CPU offloading — LRU-driven GPU→CPU eviction with on-demand prefetch.
+//!
+//! In production the `data` field of [`KvPageData`] would be a pinned-host-memory
+//! buffer obtained via `cudaMallocHost` / `hipHostMalloc` so that DMA transfers
+//! are initiated without CPU involvement.  In this simulation layer it is a plain
+//! `Vec<f32>` so the module compiles and passes on any host without a GPU driver.
+//!
+//! # Usage
+//!
+//! ```rust
+//! use rtx_inference::cache::cpu_offload::{CpuOffloadConfig, KvCpuOffloadManager};
+//! use uuid::Uuid;
+//!
+//! let cfg = CpuOffloadConfig { gpu_page_limit: 4, cpu_page_limit: 8, ..Default::default() };
+//! let mut mgr = KvCpuOffloadManager::new(cfg);
+//!
+//! let id = Uuid::new_v4();
+//! mgr.insert(id, vec![0.0f32; 128]).unwrap();
+//!
+//! let tier = mgr.access(&id).unwrap();
+//! println!("page is on {tier:?}");
+//! ```
+
+use std::collections::{HashMap, VecDeque};
+
+use super::types::{MemoryTier, PageId};
+use crate::error::{InferenceError, InferenceResult};
+
+// ---------------------------------------------------------------------------
+// Configuration
+// ---------------------------------------------------------------------------
+
+/// Configuration for CPU offload behaviour.
+#[derive(Debug, Clone)]
+pub struct CpuOffloadConfig {
+    /// GPU page capacity before offloading begins.
+    pub gpu_page_limit: usize,
+    /// Maximum pages in CPU RAM (beyond this, pages are dropped from CPU too).
+    pub cpu_page_limit: usize,
+    /// Pages per batch to offload in one DMA transfer (default: 4).
+    pub offload_batch_size: usize,
+    /// Pages per batch to prefetch back to GPU (default: 2).
+    pub prefetch_batch_size: usize,
+}
+
+impl Default for CpuOffloadConfig {
+    fn default() -> Self {
+        Self {
+            gpu_page_limit: 512,
+            cpu_page_limit: 2048,
+            offload_batch_size: 4,
+            prefetch_batch_size: 2,
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Page data
+// ---------------------------------------------------------------------------
+
+/// Simulated KV page data.
+///
+/// In production `data` would be a pinned host-memory buffer (aligned to the
+/// DMA page boundary) and `last_accessed` would be a GPU-timestamp so that
+/// pre-emptions are counted accurately across streams.
+#[derive(Debug, Clone)]
+pub struct KvPageData {
+    /// Unique page identifier.
+    pub page_id: PageId,
+    /// Current memory tier.
+    pub tier: MemoryTier,
+    /// CPU-side payload (production: pinned host memory).
+    pub data: Vec<f32>,
+    /// Logical clock tick at last access.
+    pub last_accessed: u64,
+}
+
+// ---------------------------------------------------------------------------
+// Stats snapshot
+// ---------------------------------------------------------------------------
+
+/// Snapshot of offload manager statistics.
+#[derive(Debug, Clone)]
+pub struct OffloadStats {
+    /// Pages currently on GPU.
+    pub gpu_pages: usize,
+    /// Pages currently on CPU.
+    pub cpu_pages: usize,
+    /// Total GPU→CPU offloads that have occurred.
+    pub offload_count: usize,
+    /// Total CPU→GPU prefetches that have occurred.
+    pub prefetch_count: usize,
+    /// GPU hit rate: `gpu_hits / (gpu_hits + cpu_hits)`.
+    pub gpu_hit_rate: f64,
+    /// GPU occupancy fraction (0.0–1.0).
+    pub gpu_utilization: f64,
+}
+
+// ---------------------------------------------------------------------------
+// Manager
+// ---------------------------------------------------------------------------
+
+/// Manages KV cache offloading between the GPU and CPU tiers.
+///
+/// Invariants maintained at all times:
+/// - A `PageId` appears in **at most one** of `gpu_pages` or `cpu_pages`.
+/// - `gpu_lru` is a permutation of the keys in `gpu_pages`.
+/// - `clock` is strictly monotonically increasing.
+pub struct KvCpuOffloadManager {
+    config: CpuOffloadConfig,
+    /// Pages currently on GPU: `page_id → data`.
+    gpu_pages: HashMap<PageId, KvPageData>,
+    /// Pages currently on CPU: `page_id → data`.
+    cpu_pages: HashMap<PageId, KvPageData>,
+    /// LRU order for GPU pages — front = least recently used.
+    gpu_lru: VecDeque<PageId>,
+    /// Monotonically increasing logical access clock.
+    clock: u64,
+    /// Cumulative GPU→CPU offload count.
+    offload_count: usize,
+    /// Cumulative CPU→GPU prefetch count.
+    prefetch_count: usize,
+    /// Accesses satisfied directly from GPU.
+    gpu_hits: usize,
+    /// Accesses satisfied from CPU (triggered a prefetch).
+    cpu_hits: usize,
+}
+
+impl KvCpuOffloadManager {
+    // -----------------------------------------------------------------------
+    // Construction
+    // -----------------------------------------------------------------------
+
+    /// Create a new manager with the given configuration.
+    #[must_use]
+    pub fn new(config: CpuOffloadConfig) -> Self {
+        Self {
+            config,
+            gpu_pages: HashMap::new(),
+            cpu_pages: HashMap::new(),
+            gpu_lru: VecDeque::new(),
+            clock: 0,
+            offload_count: 0,
+            prefetch_count: 0,
+            gpu_hits: 0,
+            cpu_hits: 0,
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Public API
+    // -----------------------------------------------------------------------
+
+    /// Insert a new page (allocated from GPU).
+    ///
+    /// If inserting this page would exceed [`CpuOffloadConfig::gpu_page_limit`],
+    /// a batch of the oldest GPU pages is offloaded to CPU first.
+    pub fn insert(&mut self, page_id: PageId, data: Vec<f32>) -> InferenceResult<()> {
+        // Ensure room on GPU before inserting.
+        if self.gpu_pages.len() >= self.config.gpu_page_limit {
+            let n = self.config.offload_batch_size;
+            self.offload_to_cpu(n)?;
+        }
+
+        let tick = self.next_tick();
+        let page = KvPageData {
+            page_id,
+            tier: MemoryTier::GPU,
+            data,
+            last_accessed: tick,
+        };
+
+        self.gpu_pages.insert(page_id, page);
+        self.gpu_lru.push_back(page_id);
+        Ok(())
+    }
+
+    /// Access a page by ID.
+    ///
+    /// - **GPU hit**: updates LRU position and returns [`MemoryTier::GPU`].
+    /// - **CPU hit**: prefetches the page to GPU (evicting an old GPU page if
+    ///   necessary) and returns [`MemoryTier::CPU`] to indicate it was a cold
+    ///   access.
+    /// - **Miss**: returns [`InferenceError::KvCacheError`].
+    pub fn access(&mut self, page_id: &PageId) -> InferenceResult<MemoryTier> {
+        if self.gpu_pages.contains_key(page_id) {
+            // GPU hit — refresh LRU.
+            let tick = self.next_tick();
+            let page = self
+                .gpu_pages
+                .get_mut(page_id)
+                .expect("key existence confirmed above");
+            page.last_accessed = tick;
+
+            // Move to back of LRU deque (most recently used).
+            self.gpu_lru.retain(|id| id != page_id);
+            self.gpu_lru.push_back(*page_id);
+
+            self.gpu_hits += 1;
+            return Ok(MemoryTier::GPU);
+        }
+
+        if self.cpu_pages.contains_key(page_id) {
+            // CPU hit — prefetch to GPU.
+            self.prefetch_to_gpu(&[*page_id])?;
+            self.cpu_hits += 1;
+            return Ok(MemoryTier::CPU);
+        }
+
+        Err(InferenceError::kv_cache_error(
+            "access",
+            format!("page {page_id} not found in any tier"),
+        ))
+    }
+
+    /// Evict the `n` oldest GPU pages to CPU RAM.
+    ///
+    /// If the CPU is also full, the oldest CPU page is silently dropped to make
+    /// room (simulating NVMe spillage or discard in a constrained environment).
+    ///
+    /// Returns the IDs of pages that were moved.
+    pub fn offload_to_cpu(&mut self, n: usize) -> InferenceResult<Vec<PageId>> {
+        let to_offload: Vec<PageId> = self.gpu_lru.iter().take(n).copied().collect();
+
+        if to_offload.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        let mut moved = Vec::with_capacity(to_offload.len());
+
+        for id in &to_offload {
+            if let Some(mut page) = self.gpu_pages.remove(id) {
+                // Ensure CPU has room.
+                if self.cpu_pages.len() >= self.config.cpu_page_limit {
+                    // Drop the oldest CPU page (front of insertion order is
+                    // approximated here by just removing an arbitrary entry;
+                    // a production impl would maintain a CPU LRU deque too).
+                    if let Some(victim_id) = self.cpu_pages.keys().next().copied() {
+                        self.cpu_pages.remove(&victim_id);
+                    }
+                }
+
+                let tick = self.next_tick();
+                page.tier = MemoryTier::CPU;
+                page.last_accessed = tick;
+                self.cpu_pages.insert(*id, page);
+                moved.push(*id);
+                self.offload_count += 1;
+            }
+        }
+
+        // Prune the LRU deque to match.
+        self.gpu_lru.retain(|id| !moved.contains(id));
+
+        Ok(moved)
+    }
+
+    /// Prefetch `page_ids` from CPU back to GPU.
+    ///
+    /// For each page, if the GPU is already at capacity a single LRU page is
+    /// offloaded to make room before the prefetch is executed.
+    pub fn prefetch_to_gpu(&mut self, page_ids: &[PageId]) -> InferenceResult<()> {
+        for id in page_ids {
+            if !self.cpu_pages.contains_key(id) {
+                // Already on GPU or not tracked — skip.
+                continue;
+            }
+
+            // Ensure GPU has capacity.
+            if self.gpu_pages.len() >= self.config.gpu_page_limit {
+                if let Some(victim) = self.gpu_lru.pop_front() {
+                    if let Some(mut victim_page) = self.gpu_pages.remove(&victim) {
+                        // Ensure CPU has room for the evicted GPU page.
+                        if self.cpu_pages.len() >= self.config.cpu_page_limit {
+                            if let Some(old_cpu_id) = self.cpu_pages.keys().next().copied() {
+                                self.cpu_pages.remove(&old_cpu_id);
+                            }
+                        }
+                        let tick = self.next_tick();
+                        victim_page.tier = MemoryTier::CPU;
+                        victim_page.last_accessed = tick;
+                        self.cpu_pages.insert(victim, victim_page);
+                        self.offload_count += 1;
+                    }
+                }
+            }
+
+            // Move from CPU to GPU.
+            if let Some(mut page) = self.cpu_pages.remove(id) {
+                let tick = self.next_tick();
+                page.tier = MemoryTier::GPU;
+                page.last_accessed = tick;
+                self.gpu_pages.insert(*id, page);
+                self.gpu_lru.push_back(*id);
+                self.prefetch_count += 1;
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Remove a page from all tiers (call when the request owning the page
+    /// completes and the KV data is no longer needed).
+    ///
+    /// Returns `true` if the page was found and removed, `false` otherwise.
+    pub fn evict(&mut self, page_id: &PageId) -> bool {
+        let on_gpu = self.gpu_pages.remove(page_id).is_some();
+        if on_gpu {
+            self.gpu_lru.retain(|id| id != page_id);
+            return true;
+        }
+        self.cpu_pages.remove(page_id).is_some()
+    }
+
+    /// Return `true` when the GPU page count equals or exceeds the configured
+    /// limit and the next insert would trigger an offload.
+    #[must_use]
+    pub fn gpu_pressure(&self) -> bool {
+        self.gpu_pages.len() >= self.config.gpu_page_limit
+    }
+
+    /// Return a point-in-time stats snapshot.
+    #[must_use]
+    pub fn stats(&self) -> OffloadStats {
+        let total_hits = self.gpu_hits + self.cpu_hits;
+        let gpu_hit_rate = if total_hits == 0 {
+            0.0
+        } else {
+            self.gpu_hits as f64 / total_hits as f64
+        };
+
+        OffloadStats {
+            gpu_pages: self.gpu_pages.len(),
+            cpu_pages: self.cpu_pages.len(),
+            offload_count: self.offload_count,
+            prefetch_count: self.prefetch_count,
+            gpu_hit_rate,
+            gpu_utilization: self.gpu_utilization(),
+        }
+    }
+
+    /// GPU occupancy fraction in `[0.0, 1.0]`.
+    #[must_use]
+    pub fn gpu_utilization(&self) -> f64 {
+        if self.config.gpu_page_limit == 0 {
+            return 0.0;
+        }
+        self.gpu_pages.len() as f64 / self.config.gpu_page_limit as f64
+    }
+
+    // -----------------------------------------------------------------------
+    // Private helpers
+    // -----------------------------------------------------------------------
+
+    /// Advance the logical clock and return the new tick.
+    fn next_tick(&mut self) -> u64 {
+        self.clock += 1;
+        self.clock
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use uuid::Uuid;
+
+    /// Build a small manager with deterministic limits for unit tests.
+    fn small_manager(gpu_limit: usize, cpu_limit: usize) -> KvCpuOffloadManager {
+        KvCpuOffloadManager::new(CpuOffloadConfig {
+            gpu_page_limit: gpu_limit,
+            cpu_page_limit: cpu_limit,
+            offload_batch_size: 1,
+            prefetch_batch_size: 1,
+        })
+    }
+
+    fn page_data(n: usize) -> Vec<f32> {
+        vec![n as f32; 16]
+    }
+
+    // ------------------------------------------------------------------
+    // Basic insert / access
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_insert_page_appears_on_gpu() {
+        let mut mgr = small_manager(8, 32);
+        let id = Uuid::new_v4();
+        mgr.insert(id, page_data(1)).unwrap();
+        // The page must be in gpu_pages and present via a GPU-hit access.
+        let tier = mgr.access(&id).unwrap();
+        assert_eq!(tier, MemoryTier::GPU);
+    }
+
+    #[test]
+    fn test_access_gpu_hit() {
+        let mut mgr = small_manager(8, 32);
+        let id = Uuid::new_v4();
+        mgr.insert(id, page_data(2)).unwrap();
+        let tier = mgr.access(&id).unwrap();
+        assert_eq!(tier, MemoryTier::GPU, "first access on GPU page must be a GPU hit");
+    }
+
+    #[test]
+    fn test_access_missing_page_error() {
+        let mut mgr = small_manager(8, 32);
+        let missing = Uuid::new_v4();
+        let result = mgr.access(&missing);
+        assert!(result.is_err(), "accessing an unknown page must return an error");
+        let err = result.unwrap_err();
+        assert!(
+            matches!(err, InferenceError::KvCacheError { .. }),
+            "expected KvCacheError, got {err:?}",
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // Overflow / offload triggering
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_gpu_overflow_triggers_offload() {
+        // GPU limit = 4, batch = 1.  Inserting the 5th page must push the
+        // oldest page to CPU.
+        let mut mgr = small_manager(4, 64);
+        let mut ids = Vec::new();
+        for i in 0..4 {
+            let id = Uuid::new_v4();
+            ids.push(id);
+            mgr.insert(id, page_data(i)).unwrap();
+        }
+        // GPU is now full.
+        assert_eq!(mgr.gpu_pages.len(), 4);
+
+        // Insert one more — should trigger offload of oldest (ids[0]).
+        let extra = Uuid::new_v4();
+        mgr.insert(extra, page_data(99)).unwrap();
+
+        // ids[0] must have moved to CPU.
+        assert!(
+            mgr.cpu_pages.contains_key(&ids[0]),
+            "oldest page must be on CPU after overflow",
+        );
+        // GPU still at limit or one below (we offloaded 1, inserted 1).
+        assert_eq!(mgr.gpu_pages.len(), 4);
+    }
+
+    // ------------------------------------------------------------------
+    // Explicit offload / prefetch
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_offload_to_cpu_moves_page() {
+        let mut mgr = small_manager(8, 32);
+        let id = Uuid::new_v4();
+        mgr.insert(id, page_data(10)).unwrap();
+
+        let moved = mgr.offload_to_cpu(1).unwrap();
+        assert_eq!(moved, vec![id]);
+        assert!(!mgr.gpu_pages.contains_key(&id), "page must leave gpu_pages");
+        assert!(mgr.cpu_pages.contains_key(&id), "page must enter cpu_pages");
+    }
+
+    #[test]
+    fn test_prefetch_to_gpu_restores_page() {
+        let mut mgr = small_manager(8, 32);
+        let id = Uuid::new_v4();
+        mgr.insert(id, page_data(20)).unwrap();
+        mgr.offload_to_cpu(1).unwrap();
+
+        // Page is on CPU — prefetch it back.
+        mgr.prefetch_to_gpu(&[id]).unwrap();
+        assert!(mgr.gpu_pages.contains_key(&id), "page must be back on GPU after prefetch");
+        assert!(!mgr.cpu_pages.contains_key(&id), "page must not remain on CPU");
+    }
+
+    #[test]
+    fn test_prefetch_evicts_gpu_page_when_full() {
+        // GPU limit = 2.  Fill GPU, move one page to CPU, then attempt prefetch
+        // while GPU is full — prefetch must evict an old GPU page first.
+        let mut mgr = small_manager(2, 64);
+        let id_a = Uuid::new_v4();
+        let id_b = Uuid::new_v4();
+        mgr.insert(id_a, page_data(1)).unwrap();
+        mgr.insert(id_b, page_data(2)).unwrap();
+
+        // Manually move id_a to CPU so we can try to prefetch it back while GPU
+        // is occupied by id_b and another page.
+        mgr.offload_to_cpu(1).unwrap(); // evicts id_a (oldest)
+        // GPU now has id_b (1 page).  Insert one more to fill GPU.
+        let id_c = Uuid::new_v4();
+        mgr.insert(id_c, page_data(3)).unwrap();
+        assert_eq!(mgr.gpu_pages.len(), 2, "GPU must be full before prefetch test");
+
+        // Prefetch id_a back — GPU must evict its LRU page to make room.
+        mgr.prefetch_to_gpu(&[id_a]).unwrap();
+        assert!(mgr.gpu_pages.contains_key(&id_a), "id_a must be on GPU after prefetch");
+        assert_eq!(mgr.gpu_pages.len(), 2, "GPU page count must stay at limit");
+    }
+
+    // ------------------------------------------------------------------
+    // Eviction (request-complete)
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_evict_removes_from_gpu() {
+        let mut mgr = small_manager(8, 32);
+        let id = Uuid::new_v4();
+        mgr.insert(id, page_data(5)).unwrap();
+
+        let removed = mgr.evict(&id);
+        assert!(removed, "evict must return true for a known page");
+        assert!(!mgr.gpu_pages.contains_key(&id));
+        assert!(!mgr.cpu_pages.contains_key(&id));
+    }
+
+    #[test]
+    fn test_evict_removes_from_cpu() {
+        let mut mgr = small_manager(8, 32);
+        let id = Uuid::new_v4();
+        mgr.insert(id, page_data(6)).unwrap();
+        mgr.offload_to_cpu(1).unwrap();
+
+        let removed = mgr.evict(&id);
+        assert!(removed, "evict must return true even for CPU-resident page");
+        assert!(!mgr.cpu_pages.contains_key(&id));
+    }
+
+    // ------------------------------------------------------------------
+    // Pressure / utilisation
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_gpu_pressure_true_when_full() {
+        let mut mgr = small_manager(3, 32);
+        for i in 0..3 {
+            mgr.insert(Uuid::new_v4(), page_data(i)).unwrap();
+        }
+        assert!(mgr.gpu_pressure(), "gpu_pressure must be true at limit");
+    }
+
+    #[test]
+    fn test_gpu_utilization_fraction() {
+        let mut mgr = KvCpuOffloadManager::new(CpuOffloadConfig {
+            gpu_page_limit: 512,
+            cpu_page_limit: 2048,
+            offload_batch_size: 4,
+            prefetch_batch_size: 2,
+        });
+        for _ in 0..256 {
+            mgr.insert(Uuid::new_v4(), page_data(0)).unwrap();
+        }
+        let util = mgr.gpu_utilization();
+        assert!(
+            (util - 0.5_f64).abs() < 1e-9,
+            "256/512 must give utilization ≈ 0.5, got {util}",
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // Stats
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_stats_offload_count() {
+        // Use batch_size=2: offloading 6 pages across 3 explicit calls.
+        let mut mgr = KvCpuOffloadManager::new(CpuOffloadConfig {
+            gpu_page_limit: 16,
+            cpu_page_limit: 64,
+            offload_batch_size: 2,
+            prefetch_batch_size: 2,
+        });
+        for i in 0..6 {
+            mgr.insert(Uuid::new_v4(), page_data(i)).unwrap();
+        }
+        mgr.offload_to_cpu(2).unwrap();
+        mgr.offload_to_cpu(2).unwrap();
+        mgr.offload_to_cpu(2).unwrap();
+
+        let stats = mgr.stats();
+        assert!(stats.offload_count > 0, "offload_count must be nonzero after offloading");
+        assert_eq!(stats.offload_count, 6);
+    }
+
+    // ------------------------------------------------------------------
+    // LRU ordering
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_lru_evicts_oldest_first() {
+        // Insert three pages in order; the first inserted must be offloaded
+        // when we call offload_to_cpu(1).
+        let mut mgr = small_manager(8, 32);
+        let id_first = Uuid::new_v4();
+        let id_second = Uuid::new_v4();
+        let id_third = Uuid::new_v4();
+
+        mgr.insert(id_first, page_data(1)).unwrap();
+        mgr.insert(id_second, page_data(2)).unwrap();
+        mgr.insert(id_third, page_data(3)).unwrap();
+
+        let moved = mgr.offload_to_cpu(1).unwrap();
+        assert_eq!(
+            moved[0], id_first,
+            "LRU eviction must evict the oldest-inserted page first"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // Access-from-CPU path
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_access_cpu_page_returns_cpu_tier() {
+        let mut mgr = small_manager(8, 32);
+        let id = Uuid::new_v4();
+        mgr.insert(id, page_data(42)).unwrap();
+        mgr.offload_to_cpu(1).unwrap(); // move to CPU
+
+        // Access must succeed and report the page was cold (on CPU).
+        let tier = mgr.access(&id).unwrap();
+        assert_eq!(tier, MemoryTier::CPU, "access of CPU-resident page must return Cpu tier");
+
+        // After the prefetch triggered by access, page must be on GPU now.
+        assert!(mgr.gpu_pages.contains_key(&id), "page must be on GPU after access-triggered prefetch");
+    }
+}

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

@@ -18,6 +18,7 @@
 //! sharing a common prefix (e.g. a system prompt) to skip recomputation.
 
 mod attention_sink;
+pub mod cpu_offload;
 mod eviction;
 pub mod kv_quant;
 mod manager;
@@ -29,6 +30,7 @@ mod types;
 
 // Re-export all public types
 pub use attention_sink::AttentionSinkEviction;
+pub use cpu_offload::{CpuOffloadConfig, KvCpuOffloadManager, KvPageData, OffloadStats};
 pub use eviction::AttentionScoreEviction;
 pub use kv_quant::{KvCacheQuantMode, KvQuantizer, QuantizedKvBlock};
 pub use manager::PagedKvCacheManager;

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

@@ -0,0 +1,649 @@
+//! Grouped Query Attention (GQA) KV head expansion for inference.
+//!
+//! Modern LLMs such as Llama 2/3, Mistral, Gemma, and Qwen use GQA where
+//! `num_kv_heads < num_q_heads`.  During the attention forward pass each KV
+//! head is shared by `queries_per_group = num_q_heads / num_kv_heads` query
+//! heads, eliminating the need to store a full per-head KV cache.
+//!
+//! This module provides:
+//! - [`GqaConfig`] — validated configuration object.
+//! - [`expand_kv_heads`] — repeat/tile KV tensor to match Q head count.
+//! - [`kv_head_for_q`] — zero-cost index mapping from Q head → KV head.
+//! - [`gqa_attention_cpu`] — single-batch GQA forward pass on the CPU.
+//!
+//! # Layout convention
+//!
+//! Tensors are represented as flat `&[f32]` / `Vec<f32>` in
+//! **row-major (C-contiguous)** order.  The logical shape is written in
+//! square brackets throughout the docs.
+//!
+//! ```text
+//! KV tensor: [batch, num_kv_heads, seq_len, head_dim]
+//! Q  tensor: [num_q_heads, seq_len_q, head_dim]        (single batch item)
+//! ```
+
+use thiserror::Error;
+
+// ──────────────────────────────────────────────────────────────────────────────
+// Error type
+// ──────────────────────────────────────────────────────────────────────────────
+
+/// Errors that can occur when constructing or using a [`GqaConfig`].
+#[derive(Debug, Error, PartialEq, Eq)]
+pub enum GqaError {
+    /// `num_q_heads` is not evenly divisible by `num_kv_heads`.
+    #[error("num_q_heads ({q}) must be divisible by num_kv_heads ({kv})")]
+    IndivisibleHeads { q: usize, kv: usize },
+
+    /// Either head count is zero.
+    #[error("num_q_heads and num_kv_heads must be > 0")]
+    ZeroHeads,
+
+    /// `num_kv_heads` is larger than `num_q_heads`, which is not valid.
+    #[error("num_kv_heads ({kv}) cannot exceed num_q_heads ({q})")]
+    KvExceedsQ { q: usize, kv: usize },
+
+    /// `head_dim` is zero.
+    #[error("head_dim must be > 0")]
+    ZeroHeadDim,
+}
+
+// ──────────────────────────────────────────────────────────────────────────────
+// Configuration
+// ──────────────────────────────────────────────────────────────────────────────
+
+/// Configuration for Grouped Query Attention.
+///
+/// # Invariants (upheld by [`GqaConfig::new`])
+/// - `num_q_heads > 0`
+/// - `num_kv_heads > 0`
+/// - `num_kv_heads <= num_q_heads`
+/// - `num_q_heads % num_kv_heads == 0`
+/// - `head_dim > 0`
+///
+/// # Examples
+///
+/// ```
+/// use rtx_inference::gqa::GqaConfig;
+///
+/// // Llama 3 8B: 32 Q heads, 8 KV heads, head_dim 128.
+/// let cfg = GqaConfig::new(32, 8, 128).unwrap();
+/// assert_eq!(cfg.queries_per_group(), 4);
+/// assert!(cfg.is_grouped());
+/// assert!(!cfg.is_mqa());
+/// assert!((cfg.kv_memory_ratio() - 0.25).abs() < 1e-6);
+/// ```
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct GqaConfig {
+    /// Total number of query heads.
+    pub num_q_heads: usize,
+    /// Total number of key/value heads (≤ `num_q_heads`).
+    pub num_kv_heads: usize,
+    /// Dimension of each head.
+    pub head_dim: usize,
+}
+
+impl GqaConfig {
+    /// Construct a validated [`GqaConfig`].
+    ///
+    /// # Errors
+    ///
+    /// Returns [`GqaError`] when any invariant is violated.
+    pub fn new(num_q_heads: usize, num_kv_heads: usize, head_dim: usize) -> Result<Self, GqaError> {
+        if num_q_heads == 0 || num_kv_heads == 0 {
+            return Err(GqaError::ZeroHeads);
+        }
+        if head_dim == 0 {
+            return Err(GqaError::ZeroHeadDim);
+        }
+        if num_kv_heads > num_q_heads {
+            return Err(GqaError::KvExceedsQ {
+                q: num_q_heads,
+                kv: num_kv_heads,
+            });
+        }
+        if num_q_heads % num_kv_heads != 0 {
+            return Err(GqaError::IndivisibleHeads {
+                q: num_q_heads,
+                kv: num_kv_heads,
+            });
+        }
+        Ok(Self { num_q_heads, num_kv_heads, head_dim })
+    }
+
+    /// Number of query heads that share one KV head.
+    ///
+    /// Equals `num_q_heads / num_kv_heads`.  For standard MHA this is `1`.
+    #[inline]
+    pub fn queries_per_group(&self) -> usize {
+        self.num_q_heads / self.num_kv_heads
+    }
+
+    /// Returns `true` when GQA is active (`num_kv_heads < num_q_heads`).
+    #[inline]
+    pub fn is_grouped(&self) -> bool {
+        self.num_kv_heads < self.num_q_heads
+    }
+
+    /// Returns `true` for Multi-Query Attention (single KV head for all Q heads).
+    #[inline]
+    pub fn is_mqa(&self) -> bool {
+        self.num_kv_heads == 1
+    }
+
+    /// Ratio of KV memory relative to full MHA.
+    ///
+    /// E.g. 8 KV heads / 32 Q heads = 0.25  →  75 % memory saving.
+    #[inline]
+    pub fn kv_memory_ratio(&self) -> f32 {
+        self.num_kv_heads as f32 / self.num_q_heads as f32
+    }
+}
+
+// ──────────────────────────────────────────────────────────────────────────────
+// Index mapping
+// ──────────────────────────────────────────────────────────────────────────────
+
+/// Map a query head index to its corresponding KV head index.
+///
+/// The mapping is: `kv_head = q_head / queries_per_group`.
+///
+/// This is a zero-cost inline function suitable for tight inner loops.
+///
+/// # Examples
+///
+/// ```
+/// use rtx_inference::gqa::kv_head_for_q;
+///
+/// // 8 Q heads, 2 KV heads → queries_per_group = 4
+/// assert_eq!(kv_head_for_q(0, 4), 0);
+/// assert_eq!(kv_head_for_q(3, 4), 0);
+/// assert_eq!(kv_head_for_q(4, 4), 1);
+/// assert_eq!(kv_head_for_q(7, 4), 1);
+/// ```
+#[inline]
+pub fn kv_head_for_q(q_head: usize, queries_per_group: usize) -> usize {
+    q_head / queries_per_group
+}
+
+// ──────────────────────────────────────────────────────────────────────────────
+// KV head expansion
+// ──────────────────────────────────────────────────────────────────────────────
+
+/// Expand KV heads to match the query head count by repetition (tiling).
+///
+/// Input layout:  `[batch, num_kv_heads, seq_len, head_dim]`
+/// Output layout: `[batch, num_q_heads,  seq_len, head_dim]`
+///
+/// Each KV head `i` is repeated [`GqaConfig::queries_per_group`] times
+/// consecutively so that the output can be used directly with a full-MHA
+/// kernel.
+///
+/// # Memory
+///
+/// Allocates a new `Vec<f32>` of length
+/// `batch_size * num_q_heads * seq_len * head_dim`.
+///
+/// # Panics
+///
+/// Panics in debug mode if `kv.len()` does not match the expected size
+/// `batch_size * config.num_kv_heads * seq_len * config.head_dim`.
+///
+/// # Examples
+///
+/// ```
+/// use rtx_inference::gqa::{GqaConfig, expand_kv_heads};
+///
+/// let cfg = GqaConfig::new(8, 2, 4).unwrap();
+/// // KV: [1, 2, 3, 4]  (batch=1, num_kv=2, seq=3, head_dim=4)
+/// let kv: Vec<f32> = (0..24).map(|x| x as f32).collect();
+/// let out = expand_kv_heads(&kv, &cfg, 1, 3);
+/// assert_eq!(out.len(), 1 * 8 * 3 * 4);
+/// ```
+pub fn expand_kv_heads(
+    kv: &[f32],
+    config: &GqaConfig,
+    batch_size: usize,
+    seq_len: usize,
+) -> Vec<f32> {
+    let kv_head_stride = seq_len * config.head_dim;
+    let kv_batch_stride = config.num_kv_heads * kv_head_stride;
+    let q_head_stride = seq_len * config.head_dim;
+    let q_batch_stride = config.num_q_heads * q_head_stride;
+    let gpg = config.queries_per_group();
+
+    debug_assert_eq!(
+        kv.len(),
+        batch_size * kv_batch_stride,
+        "KV slice length mismatch: expected {} got {}",
+        batch_size * kv_batch_stride,
+        kv.len()
+    );
+
+    let mut out = vec![0.0f32; batch_size * q_batch_stride];
+
+    for b in 0..batch_size {
+        for q_head in 0..config.num_q_heads {
+            let kv_head = kv_head_for_q(q_head, gpg);
+            let kv_offset = b * kv_batch_stride + kv_head * kv_head_stride;
+            let out_offset = b * q_batch_stride + q_head * q_head_stride;
+            // Copy the entire [seq_len, head_dim] slice.
+            out[out_offset..out_offset + kv_head_stride]
+                .copy_from_slice(&kv[kv_offset..kv_offset + kv_head_stride]);
+        }
+    }
+
+    out
+}
+
+// ──────────────────────────────────────────────────────────────────────────────
+// CPU attention forward pass
+// ──────────────────────────────────────────────────────────────────────────────
+
+/// Compute GQA scaled dot-product attention for a single batch item on the CPU.
+///
+/// Tensors are flat row-major slices:
+/// - `q`: `[num_q_heads, seq_len_q, head_dim]`
+/// - `k`: `[num_kv_heads, seq_len_kv, head_dim]`  *(NOT pre-expanded)*
+/// - `v`: `[num_kv_heads, seq_len_kv, head_dim]`  *(NOT pre-expanded)*
+///
+/// Returns `[num_q_heads, seq_len_q, head_dim]`.
+///
+/// For each query head `h_q`, the corresponding KV head is
+/// `h_q / queries_per_group`.  Attention is computed as:
+///
+/// ```text
+/// scores[h_q, i, j] = dot(Q[h_q, i, :], K[h_kv, j, :]) * scale
+/// attn[h_q, i, j]   = softmax(scores[h_q, i, :])[j]   (causal mask applied before softmax)
+/// out[h_q, i, :]     = sum_j attn[h_q, i, j] * V[h_kv, j, :]
+/// ```
+///
+/// When `causal = true` positions `j > i` are masked to `-∞` before softmax.
+///
+/// # Panics
+///
+/// Panics in debug mode on slice length mismatch.
+///
+/// # Examples
+///
+/// ```
+/// use rtx_inference::gqa::{GqaConfig, gqa_attention_cpu};
+///
+/// let cfg = GqaConfig::new(2, 2, 4).unwrap();
+/// let scale = (4f32).sqrt().recip();
+/// let q: Vec<f32> = (0..16).map(|x| x as f32 * 0.01).collect();
+/// let k = q.clone();
+/// let v = q.clone();
+/// let out = gqa_attention_cpu(&q, &k, &v, &cfg, 2, 2, scale, false);
+/// assert_eq!(out.len(), 2 * 2 * 4);
+/// ```
+pub fn gqa_attention_cpu(
+    q: &[f32],
+    k: &[f32],
+    v: &[f32],
+    config: &GqaConfig,
+    seq_len_q: usize,
+    seq_len_kv: usize,
+    scale: f32,
+    causal: bool,
+) -> Vec<f32> {
+    let head_dim = config.head_dim;
+    let gpg = config.queries_per_group();
+
+    debug_assert_eq!(q.len(), config.num_q_heads * seq_len_q * head_dim);
+    debug_assert_eq!(k.len(), config.num_kv_heads * seq_len_kv * head_dim);
+    debug_assert_eq!(v.len(), config.num_kv_heads * seq_len_kv * head_dim);
+
+    // Strides for input tensors.
+    let q_head_stride = seq_len_q * head_dim;
+    let kv_head_stride = seq_len_kv * head_dim;
+
+    // Output: [num_q_heads, seq_len_q, head_dim]
+    let mut output = vec![0.0f32; config.num_q_heads * seq_len_q * head_dim];
+
+    // Scratch buffer for attention scores of one query position: [seq_len_kv]
+    let mut scores = vec![0.0f32; seq_len_kv];
+
+    for h_q in 0..config.num_q_heads {
+        let h_kv = kv_head_for_q(h_q, gpg);
+        let q_base = h_q * q_head_stride;
+        let kv_base = h_kv * kv_head_stride;
+        let out_base = h_q * q_head_stride;
+
+        for i in 0..seq_len_q {
+            // Compute raw dot-product scores for query position i.
+            let q_row = &q[q_base + i * head_dim..q_base + i * head_dim + head_dim];
+
+            let max_j = if causal { i + 1 } else { seq_len_kv };
+
+            // Fill scores; positions beyond causal boundary stay at -∞.
+            for j in 0..seq_len_kv {
+                if j < max_j {
+                    let k_row = &k[kv_base + j * head_dim..kv_base + j * head_dim + head_dim];
+                    scores[j] = dot(q_row, k_row) * scale;
+                } else {
+                    scores[j] = f32::NEG_INFINITY;
+                }
+            }
+
+            // Numerically stable softmax over scores[0..seq_len_kv].
+            let max_score = scores[..seq_len_kv]
+                .iter()
+                .copied()
+                .fold(f32::NEG_INFINITY, f32::max);
+
+            let mut sum_exp = 0.0f32;
+            for s in &mut scores[..seq_len_kv] {
+                *s = (*s - max_score).exp();
+                sum_exp += *s;
+            }
+            // Guard against the degenerate all-masked case.
+            let inv_sum = if sum_exp > 0.0 { sum_exp.recip() } else { 0.0 };
+            for s in &mut scores[..seq_len_kv] {
+                *s *= inv_sum;
+            }
+
+            // Weighted sum over value vectors.
+            let out_row =
+                &mut output[out_base + i * head_dim..out_base + i * head_dim + head_dim];
+            for j in 0..seq_len_kv {
+                let attn = scores[j];
+                if attn == 0.0 {
+                    continue;
+                }
+                let v_row = &v[kv_base + j * head_dim..kv_base + j * head_dim + head_dim];
+                for d in 0..head_dim {
+                    out_row[d] += attn * v_row[d];
+                }
+            }
+        }
+    }
+
+    output
+}
+
+// ──────────────────────────────────────────────────────────────────────────────
+// Internal helpers
+// ──────────────────────────────────────────────────────────────────────────────
+
+/// Dot product of two equal-length slices.
+#[inline]
+fn dot(a: &[f32], b: &[f32]) -> f32 {
+    debug_assert_eq!(a.len(), b.len());
+    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
+}
+
+// ──────────────────────────────────────────────────────────────────────────────
+// Tests
+// ──────────────────────────────────────────────────────────────────────────────
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // ── Config validation tests ───────────────────────────────────────────────
+
+    /// MHA (num_q == num_kv) → queries_per_group == 1.
+    #[test]
+    fn test_mha_config_queries_per_group_one() {
+        let cfg = GqaConfig::new(8, 8, 64).unwrap();
+        assert_eq!(cfg.queries_per_group(), 1);
+        assert!(!cfg.is_grouped());
+    }
+
+    /// num_q not divisible by num_kv → IndivisibleHeads error.
+    #[test]
+    fn test_gqa_config_indivisible_error() {
+        let err = GqaConfig::new(7, 3, 64).unwrap_err();
+        assert_eq!(err, GqaError::IndivisibleHeads { q: 7, kv: 3 });
+    }
+
+    /// num_kv > num_q → KvExceedsQ error.
+    #[test]
+    fn test_gqa_config_kv_exceeds_q_error() {
+        let err = GqaConfig::new(4, 8, 64).unwrap_err();
+        assert_eq!(err, GqaError::KvExceedsQ { q: 4, kv: 8 });
+    }
+
+    /// num_q == 0 → ZeroHeads error.
+    #[test]
+    fn test_gqa_config_zero_heads_error() {
+        let err = GqaConfig::new(0, 0, 64).unwrap_err();
+        assert_eq!(err, GqaError::ZeroHeads);
+    }
+
+    /// num_kv == 0 with non-zero num_q → ZeroHeads error.
+    #[test]
+    fn test_gqa_config_zero_kv_heads_error() {
+        let err = GqaConfig::new(8, 0, 64).unwrap_err();
+        assert_eq!(err, GqaError::ZeroHeads);
+    }
+
+    /// num_kv < num_q → is_grouped() == true.
+    #[test]
+    fn test_is_grouped_true() {
+        let cfg = GqaConfig::new(32, 8, 128).unwrap();
+        assert!(cfg.is_grouped());
+    }
+
+    /// num_kv == 1 → is_mqa() == true.
+    #[test]
+    fn test_is_mqa() {
+        let cfg = GqaConfig::new(8, 1, 64).unwrap();
+        assert!(cfg.is_mqa());
+        assert!(cfg.is_grouped());
+    }
+
+    /// 8 kv / 32 q → kv_memory_ratio == 0.25.
+    #[test]
+    fn test_kv_memory_ratio() {
+        let cfg = GqaConfig::new(32, 8, 128).unwrap();
+        assert!((cfg.kv_memory_ratio() - 0.25_f32).abs() < 1e-6);
+    }
+
+    // ── Index mapping tests ───────────────────────────────────────────────────
+
+    /// q_head=5, queries_per_group=4 → kv_head=1.
+    #[test]
+    fn test_kv_head_for_q_mapping() {
+        assert_eq!(kv_head_for_q(5, 4), 1);
+    }
+
+    /// Full group boundary checks for queries_per_group=4.
+    #[test]
+    fn test_kv_head_for_q_boundary() {
+        // 8 Q heads, 2 KV heads, gpg=4
+        for q in 0..4usize {
+            assert_eq!(kv_head_for_q(q, 4), 0);
+        }
+        for q in 4..8usize {
+            assert_eq!(kv_head_for_q(q, 4), 1);
+        }
+    }
+
+    // ── expand_kv_heads shape test ────────────────────────────────────────────
+
+    /// Input [1, 2, 4, 8] → output [1, 8, 4, 8] for num_q=8, num_kv=2.
+    #[test]
+    fn test_expand_kv_heads_shape() {
+        let cfg = GqaConfig::new(8, 2, 8).unwrap();
+        let batch = 1;
+        let seq = 4;
+        let kv: Vec<f32> = (0..(batch * 2 * seq * 8)).map(|x| x as f32).collect();
+        let out = expand_kv_heads(&kv, &cfg, batch, seq);
+        assert_eq!(out.len(), batch * 8 * seq * 8);
+    }
+
+    /// KV[0] is repeated 4× (first four Q slots), KV[1] repeats for the rest.
+    #[test]
+    fn test_expand_kv_heads_values() {
+        // cfg: 8 Q heads, 2 KV heads, head_dim=4, batch=1, seq=1
+        let cfg = GqaConfig::new(8, 2, 4).unwrap();
+        // KV head 0: [1,2,3,4], KV head 1: [5,6,7,8]
+        let kv: Vec<f32> = vec![1., 2., 3., 4., 5., 6., 7., 8.];
+        let out = expand_kv_heads(&kv, &cfg, 1, 1);
+        // Expect first 4 Q heads to carry KV[0] and next 4 to carry KV[1].
+        for q in 0..4 {
+            let slice = &out[q * 4..(q + 1) * 4];
+            assert_eq!(slice, &[1., 2., 3., 4.], "Q head {q} should match KV head 0");
+        }
+        for q in 4..8 {
+            let slice = &out[q * 4..(q + 1) * 4];
+            assert_eq!(slice, &[5., 6., 7., 8.], "Q head {q} should match KV head 1");
+        }
+    }
+
+    // ── gqa_attention_cpu output shape ────────────────────────────────────────
+
+    /// Output shape must be [num_q_heads, seq_len_q, head_dim].
+    #[test]
+    fn test_gqa_attention_output_shape() {
+        let cfg = GqaConfig::new(4, 2, 8).unwrap();
+        let seq_q = 3;
+        let seq_kv = 5;
+        let scale = (8f32).sqrt().recip();
+        let q: Vec<f32> = vec![0.1; 4 * seq_q * 8];
+        let k: Vec<f32> = vec![0.1; 2 * seq_kv * 8];
+        let v: Vec<f32> = vec![0.2; 2 * seq_kv * 8];
+        let out = gqa_attention_cpu(&q, &k, &v, &cfg, seq_q, seq_kv, scale, false);
+        assert_eq!(out.len(), 4 * seq_q * 8);
+    }
+
+    /// When num_kv == num_q, gqa_attention_cpu must match naive MHA within 1e-5.
+    #[test]
+    fn test_gqa_attention_mha_matches_standard() {
+        let n_heads = 2;
+        let seq = 3;
+        let head_dim = 4;
+        let cfg = GqaConfig::new(n_heads, n_heads, head_dim).unwrap();
+        let scale = (head_dim as f32).sqrt().recip();
+
+        // Simple ascending values for reproducibility.
+        let q: Vec<f32> = (0..(n_heads * seq * head_dim))
+            .map(|x| x as f32 * 0.05)
+            .collect();
+        let k = q.clone();
+        let v = q.clone();
+
+        let gqa_out = gqa_attention_cpu(&q, &k, &v, &cfg, seq, seq, scale, false);
+
+        // Compute naive MHA reference (same algorithm, but explicit expansion).
+        let ref_out = naive_mha_reference(&q, &k, &v, n_heads, seq, head_dim, scale, false);
+
+        assert_eq!(gqa_out.len(), ref_out.len());
+        for (a, b) in gqa_out.iter().zip(ref_out.iter()) {
+            assert!(
+                (a - b).abs() < 1e-5,
+                "GQA/MHA mismatch: {a} vs {b}"
+            );
+        }
+    }
+
+    /// Causal mask: for seq=4, the output should reflect that future tokens are
+    /// invisible.  Specifically, position 0 may only attend to key 0, so the
+    /// output for position 0 should equal `value[0]` exactly (after softmax of
+    /// a single-element distribution = 1.0).
+    #[test]
+    fn test_gqa_attention_causal_mask() {
+        let n_heads = 1;
+        let seq = 4;
+        let head_dim = 4;
+        let cfg = GqaConfig::new(n_heads, n_heads, head_dim).unwrap();
+        let scale = (head_dim as f32).sqrt().recip();
+
+        // Distinct value vectors so we can tell which positions contributed.
+        // V row j = all j+1 (1.0, 2.0, 3.0, 4.0 for j=0,1,2,3).
+        let q: Vec<f32> = vec![0.1; n_heads * seq * head_dim];
+        let k: Vec<f32> = vec![0.0; n_heads * seq * head_dim];
+        let mut v: Vec<f32> = vec![0.0; n_heads * seq * head_dim];
+        for j in 0..seq {
+            for d in 0..head_dim {
+                v[j * head_dim + d] = (j + 1) as f32;
+            }
+        }
+
+        let out = gqa_attention_cpu(&q, &k, &v, &cfg, seq, seq, scale, true);
+
+        // Position 0 (q_i=0): can only attend to k_j=0 → attn weight = 1.0 → out = V[0] = [1,1,1,1]
+        let pos0 = &out[0..head_dim];
+        for &val in pos0 {
+            assert!(
+                (val - 1.0).abs() < 1e-5,
+                "causal pos0: expected 1.0 got {val}"
+            );
+        }
+
+        // Position 3 (q_i=3): attends to k_j ∈ {0,1,2,3} uniformly (all keys identical).
+        // Expected output = mean of V rows = (1+2+3+4)/4 = 2.5.
+        let pos3 = &out[3 * head_dim..4 * head_dim];
+        for &val in pos3 {
+            assert!(
+                (val - 2.5).abs() < 1e-4,
+                "causal pos3: expected 2.5 got {val}"
+            );
+        }
+    }
+
+    // ── Helper ────────────────────────────────────────────────────────────────
+
+    /// Naive MHA reference (full expansion then standard attention).
+    fn naive_mha_reference(
+        q: &[f32],
+        k: &[f32],
+        v: &[f32],
+        n_heads: usize,
+        seq: usize,
+        head_dim: usize,
+        scale: f32,
+        causal: bool,
+    ) -> Vec<f32> {
+        let mut out = vec![0.0f32; n_heads * seq * head_dim];
+        let head_stride = seq * head_dim;
+        let mut scores = vec![0.0f32; seq];
+
+        for h in 0..n_heads {
+            let q_base = h * head_stride;
+            let k_base = h * head_stride;
+            let v_base = h * head_stride;
+            let o_base = h * head_stride;
+
+            for i in 0..seq {
+                let q_row = &q[q_base + i * head_dim..q_base + i * head_dim + head_dim];
+                let max_j = if causal { i + 1 } else { seq };
+
+                for j in 0..seq {
+                    if j < max_j {
+                        let k_row = &k[k_base + j * head_dim..k_base + j * head_dim + head_dim];
+                        scores[j] = dot(q_row, k_row) * scale;
+                    } else {
+                        scores[j] = f32::NEG_INFINITY;
+                    }
+                }
+
+                let max_s = scores[..seq].iter().copied().fold(f32::NEG_INFINITY, f32::max);
+                let mut sum_e = 0.0f32;
+                for s in &mut scores[..seq] {
+                    *s = (*s - max_s).exp();
+                    sum_e += *s;
+                }
+                let inv = if sum_e > 0.0 { sum_e.recip() } else { 0.0 };
+                for s in &mut scores[..seq] {
+                    *s *= inv;
+                }
+
+                let o_row = &mut out[o_base + i * head_dim..o_base + i * head_dim + head_dim];
+                for j in 0..seq {
+                    let a = scores[j];
+                    if a == 0.0 {
+                        continue;
+                    }
+                    let v_row = &v[v_base + j * head_dim..v_base + j * head_dim + head_dim];
+                    for d in 0..head_dim {
+                        o_row[d] += a * v_row[d];
+                    }
+                }
+            }
+        }
+        out
+    }
+}

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

@@ -17,6 +17,8 @@
 
 pub mod batch_processor;
 pub mod cache;
+pub mod gqa;
+pub use gqa::{expand_kv_heads, gqa_attention_cpu, kv_head_for_q, GqaConfig, GqaError};
 pub mod chunked_prefill;
 pub use chunked_prefill::{ChunkedPrefillConfig, ChunkedPrefillScheduler, ChunkedStep, PrefillChunkState};
 pub mod inference_graph;
@@ -32,8 +34,9 @@ pub mod speculative;
 
 // Re-export key types for convenience
 pub use cache::{
-    AttentionScoreEviction, AttentionSinkEviction, CacheKey, CachePage, CacheStats, EvictionPolicy,
+    AttentionScoreEviction, AttentionSinkEviction, CacheKey, CachePage, CacheStats,
-    KvCacheConfig, MemoryTier, PageId, PagedKvCache, PagedKvCacheManager, PrefixIndex,
+    CpuOffloadConfig, EvictionPolicy, KvCacheConfig, KvCpuOffloadManager, MemoryTier,
+    OffloadStats, PageId, PagedKvCache, PagedKvCacheManager, PrefixIndex,
 };
 pub use engine::{
     HealthStatus, InferenceEngine, InferenceEngineConfig, MemoryStats, ModelConfig, ModelHealth,

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

@@ -12,6 +12,7 @@ pub mod polynomial_decay;
 pub mod reduce_lr_on_plateau;
 pub mod step_lr;
 pub mod warmup;
+pub mod wsd;
 
 #[cfg(all(test, feature = "disabled_tests"))]
 pub mod scheduler_integration_tests;
@@ -25,6 +26,7 @@ pub use reduce_lr_on_plateau::{
 };
 pub use step_lr::StepLRScheduler;
 pub use warmup::WarmupScheduler;
+pub use wsd::{WsdDecayType, WsdPhase, WsdScheduler};
 
 /// Trait for learning rate schedulers
 pub trait LearningRateScheduler: Send + Sync {

crates/training/rtx-transformers/src/schedulers/wsd.rs

@@ -0,0 +1,745 @@
+//! Warmup-Stable-Decay (WSD) learning rate scheduler
+//!
+//! Implements the trapezoidal LR schedule used by Mistral, MiniMax-Text-01, MegaMath,
+//! and other modern open-source LLMs. The schedule has three phases:
+//!
+//! 1. **Warmup** — linear ramp from `0` to `peak_lr` over `warmup_steps`
+//! 2. **Stable** — constant at `peak_lr` for `stable_steps` (extendable mid-training)
+//! 3. **Decay** — cosine, linear, or sqrt decay from `peak_lr` to `min_lr` over `decay_steps`
+//!
+//! # Key property
+//!
+//! `stable_steps` can be extended at any point during training (via [`WsdScheduler::extend_stable`])
+//! without resetting the scheduler or restarting training. This enables flexible compute budgets
+//! where you decide when to stop spending tokens on the plateau.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::schedulers::wsd::{WsdScheduler, WsdDecayType};
+//! use rtx_transformers::schedulers::LearningRateScheduler;
+//!
+//! let mut sched = WsdScheduler::cosine(3e-4, 1e-5, 100, 900, 200).unwrap();
+//!
+//! // Warmup: LR ramps up
+//! let lr_step0 = sched.get_lr(0, 0);
+//! assert!(lr_step0 < 3e-4);
+//!
+//! // Stable: LR stays at peak
+//! let lr_stable = sched.get_lr(0, 500);
+//! assert!((lr_stable - 3e-4).abs() < 1e-12);
+//!
+//! // Extend stable phase by 500 more steps mid-run
+//! sched.extend_stable(500);
+//! assert_eq!(sched.total_steps(), 100 + 1400 + 200);
+//! ```
+
+use std::f64::consts::PI;
+
+use crate::schedulers::LearningRateScheduler;
+use crate::{Result, TransformerError};
+use serde::{Deserialize, Serialize};
+use tracing::{debug, trace};
+
+// ---------------------------------------------------------------------------
+// Decay variant
+// ---------------------------------------------------------------------------
+
+/// Decay function applied during the third (decay) phase.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub enum WsdDecayType {
+    /// Cosine decay from `peak_lr` to `min_lr` (smoother, generally recommended).
+    ///
+    /// `lr(p) = min_lr + (peak_lr - min_lr) * 0.5 * (1 + cos(π * p))`
+    Cosine,
+    /// Linear decay from `peak_lr` to `min_lr`.
+    ///
+    /// `lr(p) = peak_lr + (min_lr - peak_lr) * p`
+    Linear,
+    /// Square-root decay from `peak_lr` to `min_lr`.
+    ///
+    /// `lr(p) = min_lr + (peak_lr - min_lr) * (1 - sqrt(p))`
+    Sqrt,
+}
+
+// ---------------------------------------------------------------------------
+// Phase enum
+// ---------------------------------------------------------------------------
+
+/// Phase the scheduler is currently in, for external introspection.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum WsdPhase {
+    /// Linear warmup — LR is still ramping up.
+    Warmup,
+    /// Plateau — LR is constant at `peak_lr`.
+    Stable,
+    /// Decay — LR is decreasing toward `min_lr`.
+    Decay,
+    /// Past the end of the schedule — LR is clamped to `min_lr`.
+    Complete,
+}
+
+// ---------------------------------------------------------------------------
+// Scheduler struct
+// ---------------------------------------------------------------------------
+
+/// Warmup-Stable-Decay learning rate scheduler.
+///
+/// Three-phase schedule: linear warmup → constant plateau → cosine/linear/sqrt decay.
+///
+/// See the [module-level documentation](self) for a full description and usage example.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct WsdScheduler {
+    /// Target learning rate at the peak (end of warmup / entire stable phase).
+    pub peak_lr: f64,
+    /// Minimum learning rate at the end of the decay phase (and beyond).
+    pub min_lr: f64,
+    /// Number of steps in the linear warmup phase.
+    pub warmup_steps: usize,
+    /// Number of steps in the stable (constant) phase. Can be extended at any time.
+    pub stable_steps: usize,
+    /// Number of steps in the decay phase.
+    pub decay_steps: usize,
+    /// Which decay function to use during the third phase.
+    pub decay_type: WsdDecayType,
+    /// Internal step counter; advanced by [`Self::step`].
+    current_step: usize,
+}
+
+impl WsdScheduler {
+    // ------------------------------------------------------------------
+    // Constructors
+    // ------------------------------------------------------------------
+
+    /// Create a new WSD scheduler with explicit parameters.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`TransformerError`] when:
+    /// - `peak_lr` is not positive, or `peak_lr <= min_lr`
+    /// - `min_lr` is negative
+    /// - `warmup_steps` is zero
+    /// - `decay_steps` is zero
+    pub fn new(
+        peak_lr: f64,
+        min_lr: f64,
+        warmup_steps: usize,
+        stable_steps: usize,
+        decay_steps: usize,
+        decay_type: WsdDecayType,
+    ) -> Result<Self> {
+        if peak_lr <= 0.0 {
+            return Err(TransformerError::generic(format!(
+                "peak_lr {peak_lr} must be positive"
+            )));
+        }
+        if min_lr < 0.0 {
+            return Err(TransformerError::generic(format!(
+                "min_lr {min_lr} must be non-negative"
+            )));
+        }
+        if peak_lr <= min_lr {
+            return Err(TransformerError::generic(format!(
+                "peak_lr {peak_lr} must be greater than min_lr {min_lr}"
+            )));
+        }
+        if warmup_steps == 0 {
+            return Err(TransformerError::generic(
+                "warmup_steps must be greater than 0".to_string(),
+            ));
+        }
+        if decay_steps == 0 {
+            return Err(TransformerError::generic(
+                "decay_steps must be greater than 0".to_string(),
+            ));
+        }
+
+        debug!(
+            "Creating WsdScheduler: peak_lr={peak_lr}, min_lr={min_lr}, \
+             warmup={warmup_steps}, stable={stable_steps}, decay={decay_steps}, \
+             decay_type={decay_type:?}"
+        );
+
+        Ok(Self {
+            peak_lr,
+            min_lr,
+            warmup_steps,
+            stable_steps,
+            decay_steps,
+            decay_type,
+            current_step: 0,
+        })
+    }
+
+    /// Convenience constructor that uses cosine decay (the most common choice).
+    ///
+    /// # Errors
+    ///
+    /// Propagates validation errors from [`Self::new`].
+    pub fn cosine(
+        peak_lr: f64,
+        min_lr: f64,
+        warmup_steps: usize,
+        stable_steps: usize,
+        decay_steps: usize,
+    ) -> Result<Self> {
+        Self::new(
+            peak_lr,
+            min_lr,
+            warmup_steps,
+            stable_steps,
+            decay_steps,
+            WsdDecayType::Cosine,
+        )
+    }
+
+    // ------------------------------------------------------------------
+    // Mutation helpers
+    // ------------------------------------------------------------------
+
+    /// Extend the stable phase by `extra_steps` without resetting the scheduler.
+    ///
+    /// This is the key flexibility property of WSD: you can decide mid-training
+    /// to spend more compute at the peak learning rate before decaying.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use rtx_transformers::schedulers::wsd::WsdScheduler;
+    ///
+    /// let mut sched = WsdScheduler::cosine(1e-3, 1e-5, 100, 500, 100).unwrap();
+    /// assert_eq!(sched.total_steps(), 700);
+    /// sched.extend_stable(200);
+    /// assert_eq!(sched.total_steps(), 900);
+    /// ```
+    pub fn extend_stable(&mut self, extra_steps: usize) {
+        self.stable_steps = self.stable_steps.saturating_add(extra_steps);
+        debug!("Extended stable phase by {extra_steps}; new stable_steps={}", self.stable_steps);
+    }
+
+    // ------------------------------------------------------------------
+    // Queries
+    // ------------------------------------------------------------------
+
+    /// Total steps across all three phases.
+    #[must_use]
+    pub fn total_steps(&self) -> usize {
+        self.warmup_steps
+            .saturating_add(self.stable_steps)
+            .saturating_add(self.decay_steps)
+    }
+
+    /// Which phase the scheduler is in at `step`.
+    ///
+    /// Uses the *configured* `stable_steps` (including any extensions).
+    #[must_use]
+    pub fn phase_at(&self, step: usize) -> WsdPhase {
+        if step < self.warmup_steps {
+            WsdPhase::Warmup
+        } else if step < self.warmup_steps.saturating_add(self.stable_steps) {
+            WsdPhase::Stable
+        } else if step < self.total_steps() {
+            WsdPhase::Decay
+        } else {
+            WsdPhase::Complete
+        }
+    }
+
+    /// Which phase the scheduler is currently in (based on `current_step`).
+    #[must_use]
+    pub fn current_phase(&self) -> WsdPhase {
+        self.phase_at(self.current_step)
+    }
+
+    /// LR at the current internal step.
+    #[must_use]
+    pub fn current_lr(&self) -> f64 {
+        self.get_lr(0, self.current_step)
+    }
+
+    /// Fraction through the decay phase at `step`, in `0.0..=1.0`.
+    ///
+    /// Returns `0.0` for steps outside the decay phase.
+    #[must_use]
+    pub fn decay_progress_at(&self, step: usize) -> f64 {
+        let decay_start = self.warmup_steps.saturating_add(self.stable_steps);
+        if step < decay_start || step >= self.total_steps() {
+            0.0
+        } else {
+            let t = step - decay_start;
+            // Clamp to [0, 1] — last decay step maps to exactly 1.0
+            (t as f64 / self.decay_steps as f64).min(1.0)
+        }
+    }
+
+    /// Fraction through the decay phase at `current_step`.
+    #[must_use]
+    pub fn decay_progress(&self) -> f64 {
+        self.decay_progress_at(self.current_step)
+    }
+
+    // ------------------------------------------------------------------
+    // Core LR computation (pure, no mutation)
+    // ------------------------------------------------------------------
+
+    fn compute_lr(&self, step: usize) -> f64 {
+        let decay_start = self.warmup_steps.saturating_add(self.stable_steps);
+
+        if step < self.warmup_steps {
+            // Phase 1: linear warmup — lr = peak_lr * (step + 1) / warmup_steps
+            // Using (step + 1) so step=0 gives a non-zero but tiny LR, and
+            // step = warmup_steps - 1 gives peak_lr * (warmup_steps / warmup_steps) = peak_lr.
+            let progress = (step + 1) as f64 / self.warmup_steps as f64;
+            let lr = self.peak_lr * progress;
+            trace!("WSD warmup  step={step}: progress={progress:.4}, lr={lr:.6e}");
+            lr
+        } else if step < decay_start {
+            // Phase 2: stable plateau
+            trace!("WSD stable  step={step}: lr={:.6e}", self.peak_lr);
+            self.peak_lr
+        } else if step < self.total_steps() {
+            // Phase 3: decay
+            let t = step - decay_start;
+            let p = (t as f64 / self.decay_steps as f64).clamp(0.0, 1.0);
+            let lr = match self.decay_type {
+                WsdDecayType::Cosine => {
+                    self.min_lr + (self.peak_lr - self.min_lr) * 0.5 * (1.0 + (PI * p).cos())
+                }
+                WsdDecayType::Linear => self.peak_lr + (self.min_lr - self.peak_lr) * p,
+                WsdDecayType::Sqrt => {
+                    self.min_lr + (self.peak_lr - self.min_lr) * (1.0 - p.sqrt())
+                }
+            };
+            trace!("WSD decay   step={step}: p={p:.4}, lr={lr:.6e}");
+            lr
+        } else {
+            // Phase 4: complete — clamp to min_lr
+            trace!("WSD complete step={step}: lr={:.6e}", self.min_lr);
+            self.min_lr
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Trait implementation
+// ---------------------------------------------------------------------------
+
+impl LearningRateScheduler for WsdScheduler {
+    /// Compute LR at the given `step` (the `epoch` parameter is ignored — WSD
+    /// is step-based, consistent with how it is used in large LLM training).
+    fn get_lr(&self, _epoch: usize, step: usize) -> f64 {
+        self.compute_lr(step)
+    }
+
+    /// Advance the internal step counter by one.
+    fn step(&mut self) {
+        self.current_step += 1;
+        trace!("WSD scheduler stepped to {}", self.current_step);
+    }
+
+    fn current_step(&self) -> usize {
+        self.current_step
+    }
+
+    fn reset(&mut self) {
+        self.current_step = 0;
+        debug!("Reset WsdScheduler");
+    }
+
+    fn scheduler_type(&self) -> &'static str {
+        "WarmupStableDecay"
+    }
+
+    fn base_lr(&self) -> f64 {
+        self.peak_lr
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const PEAK: f64 = 3e-4;
+    const MIN: f64 = 1e-5;
+    const WARMUP: usize = 100;
+    const STABLE: usize = 900;
+    const DECAY: usize = 200;
+
+    fn default_sched() -> WsdScheduler {
+        WsdScheduler::cosine(PEAK, MIN, WARMUP, STABLE, DECAY).unwrap()
+    }
+
+    // ------------------------------------------------------------------
+    // Warmup phase
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_warmup_phase_starts_at_zero() {
+        let s = default_sched();
+        // step=0 gives peak_lr * 1/100, which is very small but non-zero
+        let lr0 = s.get_lr(0, 0);
+        assert!(lr0 > 0.0, "lr at step 0 should be > 0");
+        assert!(lr0 < PEAK, "lr at step 0 should be less than peak_lr");
+
+        // step warmup_steps-1 should be strictly less than peak_lr
+        let lr_last_warmup = s.get_lr(0, WARMUP - 1);
+        // (100/100) * PEAK = PEAK — actually at WARMUP-1 = 99: progress = 100/100 = 1.0 → PEAK
+        // Wait: step=99, (99+1)/100 = 100/100 = 1.0 → lr = PEAK
+        // So last warmup step IS peak. Let's check step 98 instead.
+        let lr_98 = s.get_lr(0, 98);
+        assert!(lr_98 < PEAK);
+        // At step 99 (last warmup step), progress = 100/100 = 1.0, so lr = PEAK
+        assert!((lr_last_warmup - PEAK).abs() < 1e-12);
+    }
+
+    #[test]
+    fn test_warmup_phase_reaches_peak() {
+        let s = default_sched();
+        // First step of stable phase = warmup_steps (step == WARMUP)
+        let lr = s.get_lr(0, WARMUP);
+        assert!(
+            (lr - PEAK).abs() < 1e-12,
+            "LR at step=warmup_steps must equal peak_lr, got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_warmup_linear_increase() {
+        let s = default_sched();
+        // LR must be monotonically increasing during warmup
+        let mut prev = s.get_lr(0, 0);
+        for step in 1..WARMUP {
+            let cur = s.get_lr(0, step);
+            assert!(
+                cur >= prev,
+                "warmup should be monotonically non-decreasing: step {step} lr={cur} < prev={prev}"
+            );
+            prev = cur;
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // Stable phase
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_stable_phase_constant() {
+        let s = default_sched();
+        for step in WARMUP..=(WARMUP + STABLE - 1) {
+            let lr = s.get_lr(0, step);
+            assert!(
+                (lr - PEAK).abs() < 1e-12,
+                "stable phase step {step}: expected {PEAK} got {lr}"
+            );
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // Decay phase — cosine
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_decay_cosine_at_start() {
+        let s = default_sched();
+        // First decay step: p = 1/200 = 0.005; cos(π * 0.005) is very close to 1.
+        // lr ≈ peak_lr — the cosine has barely moved away from the peak.
+        let decay_start = WARMUP + STABLE;
+        let lr = s.get_lr(0, decay_start);
+        assert!(
+            (lr - PEAK).abs() < 1e-3,
+            "first decay step should be ≈ peak_lr (within 1e-3), got {lr}"
+        );
+        // Must be less than or equal to peak_lr — the decay cannot overshoot.
+        assert!(
+            lr <= PEAK,
+            "first decay step must not exceed peak_lr, got {lr}"
+        );
+        // And it must be greater than min_lr — we are only at the very start.
+        assert!(
+            lr > MIN,
+            "first decay step must still be far above min_lr, got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_decay_cosine_at_end() {
+        let s = default_sched();
+        // Last decay step: step = total_steps - 1, p = (DECAY-1)/DECAY
+        let last = s.total_steps() - 1;
+        let lr = s.get_lr(0, last);
+        // cos(π * (199/200)) ≈ -1 → lr ≈ min_lr
+        assert!(
+            (lr - MIN).abs() < 1e-3,
+            "last decay step should be ≈ min_lr, got {lr}"
+        );
+        assert!(lr > MIN, "last decay step should be strictly greater than min_lr");
+    }
+
+    #[test]
+    fn test_decay_cosine_monotone_decreasing() {
+        let s = default_sched();
+        let decay_start = WARMUP + STABLE;
+        let mut prev = s.get_lr(0, decay_start);
+        for step in (decay_start + 1)..s.total_steps() {
+            let cur = s.get_lr(0, step);
+            assert!(
+                cur <= prev + 1e-15,
+                "cosine decay should be monotone non-increasing: step {step} lr={cur} > prev={prev}"
+            );
+            prev = cur;
+        }
+    }
+
+    #[test]
+    fn test_decay_linear_midpoint() {
+        let mut s = WsdScheduler::new(PEAK, MIN, WARMUP, STABLE, DECAY, WsdDecayType::Linear)
+            .unwrap();
+        // extend stable to 0 so midpoint is clearly mid-decay
+        s.stable_steps = 0;
+        let decay_start = WARMUP;
+        let mid = decay_start + DECAY / 2;
+        let lr = s.get_lr(0, mid);
+        let expected = PEAK + (MIN - PEAK) * 0.5;
+        assert!(
+            (lr - expected).abs() < 1e-10,
+            "linear decay midpoint: expected {expected}, got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_decay_sqrt_midpoint() {
+        let s =
+            WsdScheduler::new(PEAK, MIN, WARMUP, STABLE, DECAY, WsdDecayType::Sqrt).unwrap();
+        let decay_start = WARMUP + STABLE;
+        let mid = decay_start + DECAY / 2;
+        let p = 0.5_f64;
+        let lr = s.get_lr(0, mid);
+        let expected = MIN + (PEAK - MIN) * (1.0 - p.sqrt());
+        assert!(
+            (lr - expected).abs() < 1e-8,
+            "sqrt decay midpoint: expected {expected}, got {lr}"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // Complete phase
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_complete_phase_returns_min_lr() {
+        let s = default_sched();
+        let total = s.total_steps();
+        for step in [total, total + 1, total + 1000] {
+            let lr = s.get_lr(0, step);
+            assert!(
+                (lr - MIN).abs() < 1e-12,
+                "step {step} (past end) should return min_lr={MIN}, got {lr}"
+            );
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // total_steps
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_total_steps() {
+        let s = default_sched();
+        assert_eq!(
+            s.total_steps(),
+            WARMUP + STABLE + DECAY,
+            "total_steps() must equal warmup + stable + decay"
+        );
+    }
+
+    // ------------------------------------------------------------------
+    // Phase introspection
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_current_phase_warmup() {
+        let mut s = default_sched();
+        // current_step starts at 0 → Warmup
+        assert_eq!(s.current_phase(), WsdPhase::Warmup);
+        assert_eq!(s.phase_at(0), WsdPhase::Warmup);
+        assert_eq!(s.phase_at(WARMUP - 1), WsdPhase::Warmup);
+        // Step past warmup
+        for _ in 0..WARMUP {
+            s.step();
+        }
+        assert_ne!(s.current_phase(), WsdPhase::Warmup);
+    }
+
+    #[test]
+    fn test_current_phase_stable() {
+        let s = default_sched();
+        assert_eq!(s.phase_at(WARMUP), WsdPhase::Stable);
+        assert_eq!(s.phase_at(WARMUP + STABLE - 1), WsdPhase::Stable);
+    }
+
+    #[test]
+    fn test_current_phase_decay() {
+        let s = default_sched();
+        let decay_start = WARMUP + STABLE;
+        assert_eq!(s.phase_at(decay_start), WsdPhase::Decay);
+        assert_eq!(s.phase_at(decay_start + DECAY - 1), WsdPhase::Decay);
+    }
+
+    #[test]
+    fn test_current_phase_complete() {
+        let s = default_sched();
+        assert_eq!(s.phase_at(s.total_steps()), WsdPhase::Complete);
+        assert_eq!(s.phase_at(s.total_steps() + 999), WsdPhase::Complete);
+    }
+
+    // ------------------------------------------------------------------
+    // extend_stable
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_extend_stable_delays_decay() {
+        let mut s = default_sched();
+        let original_total = s.total_steps();
+        s.extend_stable(100);
+        assert_eq!(s.stable_steps, STABLE + 100);
+        assert_eq!(s.total_steps(), original_total + 100);
+
+        // The step that used to be the first decay step is now still stable
+        let old_decay_start = WARMUP + STABLE;
+        assert_eq!(
+            s.phase_at(old_decay_start),
+            WsdPhase::Stable,
+            "old decay start should now be in Stable after extension"
+        );
+        // New decay start
+        let new_decay_start = WARMUP + STABLE + 100;
+        assert_eq!(s.phase_at(new_decay_start), WsdPhase::Decay);
+    }
+
+    #[test]
+    fn test_extend_stable_lr_still_peak_at_extended_steps() {
+        let mut s = default_sched();
+        s.extend_stable(300);
+        // Any step in [WARMUP, WARMUP + STABLE + 300) should return PEAK
+        for step in [WARMUP, WARMUP + STABLE, WARMUP + STABLE + 299] {
+            let lr = s.get_lr(0, step);
+            assert!(
+                (lr - PEAK).abs() < 1e-12,
+                "extended stable step {step} should return peak_lr, got {lr}"
+            );
+        }
+    }
+
+    // ------------------------------------------------------------------
+    // decay_progress
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_decay_progress_zero_in_stable() {
+        let s = default_sched();
+        // During warmup and stable, decay_progress should be 0.0
+        for step in [0, WARMUP / 2, WARMUP, WARMUP + STABLE - 1] {
+            assert_eq!(
+                s.decay_progress_at(step),
+                0.0,
+                "step {step} is not in decay, progress should be 0.0"
+            );
+        }
+    }
+
+    #[test]
+    fn test_decay_progress_one_at_end() {
+        let s = default_sched();
+        // Last decay step: step = total_steps - 1
+        // t = total_steps - 1 - (warmup + stable) = DECAY - 1
+        // p = (DECAY - 1) / DECAY
+        let last_decay = s.total_steps() - 1;
+        let p = s.decay_progress_at(last_decay);
+        let expected = (DECAY - 1) as f64 / DECAY as f64;
+        assert!(
+            (p - expected).abs() < 1e-12,
+            "decay progress at last decay step: expected {expected}, got {p}"
+        );
+        // And complete phase (past total_steps) gives 0.0 (outside decay)
+        assert_eq!(s.decay_progress_at(s.total_steps()), 0.0);
+    }
+
+    // ------------------------------------------------------------------
+    // Trait method wiring
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_step_advances_current_step() {
+        let mut s = default_sched();
+        assert_eq!(s.current_step(), 0);
+        s.step();
+        assert_eq!(s.current_step(), 1);
+        s.step();
+        assert_eq!(s.current_step(), 2);
+    }
+
+    #[test]
+    fn test_reset_returns_to_zero() {
+        let mut s = default_sched();
+        for _ in 0..500 {
+            s.step();
+        }
+        assert_eq!(s.current_step(), 500);
+        s.reset();
+        assert_eq!(s.current_step(), 0);
+    }
+
+    #[test]
+    fn test_scheduler_type_str() {
+        let s = default_sched();
+        assert_eq!(s.scheduler_type(), "WarmupStableDecay");
+    }
+
+    #[test]
+    fn test_base_lr_returns_peak() {
+        let s = default_sched();
+        assert_eq!(s.base_lr(), PEAK);
+    }
+
+    #[test]
+    fn test_current_lr_matches_get_lr() {
+        let mut s = default_sched();
+        for _ in 0..750 {
+            s.step();
+        }
+        assert!((s.current_lr() - s.get_lr(0, s.current_step())).abs() < 1e-15);
+    }
+
+    // ------------------------------------------------------------------
+    // Validation
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_invalid_params_rejected() {
+        // negative peak_lr
+        assert!(WsdScheduler::cosine(-1e-3, 1e-5, 100, 100, 100).is_err());
+        // zero peak_lr
+        assert!(WsdScheduler::cosine(0.0, 1e-5, 100, 100, 100).is_err());
+        // peak_lr <= min_lr
+        assert!(WsdScheduler::cosine(1e-5, 1e-5, 100, 100, 100).is_err());
+        assert!(WsdScheduler::cosine(1e-6, 1e-5, 100, 100, 100).is_err());
+        // negative min_lr
+        assert!(WsdScheduler::cosine(1e-3, -1.0, 100, 100, 100).is_err());
+        // zero warmup_steps
+        assert!(WsdScheduler::cosine(1e-3, 1e-5, 0, 100, 100).is_err());
+        // zero decay_steps
+        assert!(WsdScheduler::cosine(1e-3, 1e-5, 100, 100, 0).is_err());
+    }
+
+    #[test]
+    fn test_zero_stable_steps_allowed() {
+        // stable_steps = 0 is valid (warmup directly into decay)
+        let s = WsdScheduler::cosine(PEAK, MIN, WARMUP, 0, DECAY).unwrap();
+        assert_eq!(s.total_steps(), WARMUP + DECAY);
+        assert_eq!(s.phase_at(WARMUP), WsdPhase::Decay);
+    }
+}