feat(perf): GPU perf batch 2 — EAGLE-3 dynamic draft trees + GaLore-2 optimizer

EAGLE-3 dynamic draft trees (rtx-inference)
- `eagle3.rs`: `Eagle3Config` (min_depth=1, max_depth=6, expansion_threshold=0.4,
  beam_width=3, self_consistency=true, prune_threshold=0.05), `DynamicDraftTree` with
  confidence-gated BFS expansion + iterative bottom-up cascade pruning + `all_paths()` /
  `accept_path()`, `Eagle3Decoder::build_draft_tree()` with cheap hidden-state proxy
  for child nodes (parent states scaled by child probability)
- `tree.rs`: added `path_probability(leaf)`, `leaves()` (tombstone-safe DFS)
- `types.rs`: added `DraftModelType::Eagle3` variant
- 10 new unit tests via `FixedProbDraftModel` mock (no GPU required); total 85 pass

GaLore-2 low-rank optimizer state (rtx-transformers)
- `galore.rs`: `GaLoreConfig` (rank=128, update_proj_gap=200, scale=0.25,
  min_param_size=4096, momentum_inheritance=true), `GaLoreParamState` (proj_matrix
  [rows×rank], m_lr/v_lr [rank×cols]), `GaLoreAdamW` implementing `Optimizer` trait
- Randomized range-finder sketched SVD: Ω~N(0,1) via LCG, Y=G@Ω, Gram-Schmidt QR
- Momentum inheritance: project old m_lr onto new subspace on refresh
- Automatic fallback to standard AdamW for params smaller than `min_param_size`
- Memory ratio at rank=64, param=256×256: 2×(64×256) vs 2×(256²) = 25% of full state
- `mod.rs`: `pub mod galore` + re-exports
- 12 unit tests (all CPU); total 102+12 pass

Combined: 85 + 114 = 199 lib tests pass across rtx-inference and rtx-transformers

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

crates/production/rtx-inference/src/speculative/eagle3.rs

@@ -0,0 +1,709 @@
+//! EAGLE-3 dynamic draft tree for speculative decoding.
+//!
+//! Implements the algorithm described in arXiv:2503.01840.  The key idea is
+//! that the draft tree is *grown dynamically*: nodes whose token probability
+//! exceeds `expansion_threshold` are expanded to deeper levels, while nodes
+//! below that threshold stay as leaves.  After the tree is built, an optional
+//! self-consistency re-scoring pass prunes branches whose cumulative path
+//! probability falls below `prune_threshold`.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_inference::speculative::eagle3::{Eagle3Config, Eagle3Decoder, DynamicDraftTree};
+//!
+//! let config = Eagle3Config::default();
+//! assert_eq!(config.min_depth, 1);
+//! assert_eq!(config.max_depth, 6);
+//! ```
+
+use super::eagle::{EagleConfig, EagleDraftModel};
+use super::tree::CandidateTree;
+use crate::InferenceResult;
+
+// ---------------------------------------------------------------------------
+// Configuration
+// ---------------------------------------------------------------------------
+
+/// Configuration for EAGLE-3 dynamic draft tree generation.
+#[derive(Debug, Clone)]
+pub struct Eagle3Config {
+    /// Minimum draft tree depth — always expand at least this many levels.
+    pub min_depth: usize,
+    /// Maximum draft tree depth — never expand beyond this level.
+    pub max_depth: usize,
+    /// Confidence threshold in `[0, 1]`.  A node is eligible for deeper
+    /// expansion only when its token probability exceeds this value.
+    pub expansion_threshold: f32,
+    /// Maximum number of children (beam width) per node.  Only the top-k
+    /// candidates by probability are kept at each expansion step.
+    pub beam_width: usize,
+    /// When `true`, a self-consistency re-scoring pass prunes branches whose
+    /// cumulative path probability is below `prune_threshold`.
+    pub self_consistency: bool,
+    /// Minimum path probability to survive self-consistency pruning.
+    pub prune_threshold: f32,
+    /// Sampling temperature applied to draft logits (≥ 0; 0 = greedy).
+    pub temperature: f32,
+}
+
+impl Default for Eagle3Config {
+    fn default() -> Self {
+        Self {
+            min_depth: 1,
+            max_depth: 6,
+            expansion_threshold: 0.4,
+            beam_width: 3,
+            self_consistency: true,
+            prune_threshold: 0.05,
+            temperature: 1.0,
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// DynamicDraftTree
+// ---------------------------------------------------------------------------
+
+/// A draft tree that can be grown node-by-node at runtime.
+///
+/// Wraps [`CandidateTree`] and adds the expansion / pruning logic needed by
+/// EAGLE-3.  All indices exposed by this type are *node indices* into the
+/// underlying `CandidateTree::nodes` vec.
+#[derive(Debug, Clone)]
+pub struct DynamicDraftTree {
+    /// Underlying static tree storage.
+    pub inner: CandidateTree,
+    /// Config snapshot used when the tree was constructed.
+    config: Eagle3Config,
+}
+
+impl DynamicDraftTree {
+    /// Create an empty tree with the given EAGLE-3 config.
+    #[must_use]
+    pub fn new(config: Eagle3Config) -> Self {
+        Self {
+            inner: CandidateTree::new(),
+            config,
+        }
+    }
+
+    /// Add root-level candidates (depth 0).
+    ///
+    /// Returns the indices of the newly added root nodes.
+    pub fn add_roots(&mut self, candidates: &[(u32, f32)]) -> Vec<usize> {
+        let sorted = Self::top_k_sorted(candidates, self.config.beam_width);
+        sorted
+            .iter()
+            .map(|&(token_id, prob)| self.inner.add_root(token_id, prob))
+            .collect()
+    }
+
+    /// Add children to an existing node, keeping at most `beam_width` children
+    /// sorted by probability (highest first).
+    ///
+    /// Returns the indices of the newly created child nodes.
+    pub fn expand_node(&mut self, parent_idx: usize, candidates: &[(u32, f32)]) -> Vec<usize> {
+        let sorted = Self::top_k_sorted(candidates, self.config.beam_width);
+        sorted
+            .iter()
+            .map(|&(token_id, prob)| self.inner.add_child(parent_idx, token_id, prob))
+            .collect()
+    }
+
+    /// Returns `true` when a node is eligible for deeper expansion:
+    /// - its probability exceeds `expansion_threshold`, AND
+    /// - the node's depth is less than `max_depth - 1`.
+    ///
+    /// The depth check ensures that children created from this node will be at
+    /// depth ≤ `max_depth - 1`, which makes each root-to-leaf path exactly
+    /// `max_depth` tokens long (root at depth 0, deepest leaf at depth
+    /// `max_depth - 1`).
+    ///
+    /// Panics in debug mode if `node_idx` is out of range.
+    #[must_use]
+    pub fn should_expand(&self, node_idx: usize) -> bool {
+        let node = &self.inner.nodes[node_idx];
+        // Guard: max_depth must be at least 1 for expansion to make sense.
+        if self.config.max_depth == 0 {
+            return false;
+        }
+        node.probability > self.config.expansion_threshold
+            && node.depth < self.config.max_depth - 1
+    }
+
+    /// Remove all leaf nodes whose cumulative path probability (root→leaf
+    /// product) is strictly below `threshold`, then repeat until no further
+    /// pruning is possible (bottom-up wave).
+    ///
+    /// Internal nodes that become leaves after their only child is pruned are
+    /// pruned in subsequent passes.  The loop terminates in at most `O(depth)`
+    /// iterations.
+    ///
+    /// Pruned nodes are left as tombstones in `CandidateTree::nodes` (their
+    /// `children` list is cleared and they are removed from their parent's
+    /// `children` list / `roots` vec).  Because `get_all_paths` only visits
+    /// nodes reachable from `roots`, tombstones are never included in paths.
+    pub fn prune_by_score(&mut self, threshold: f32) {
+        loop {
+            let to_remove: Vec<usize> = self
+                .inner
+                .leaves()
+                .into_iter()
+                .filter(|&idx| self.inner.path_probability(idx) < threshold)
+                .collect();
+
+            if to_remove.is_empty() {
+                break;
+            }
+
+            // Detach each pruned leaf from its parent / roots list.
+            for &leaf_idx in &to_remove {
+                if let Some(parent_idx) = self.inner.nodes[leaf_idx].parent {
+                    self.inner.nodes[parent_idx]
+                        .children
+                        .retain(|&c| c != leaf_idx);
+                } else {
+                    self.inner.roots.retain(|&r| r != leaf_idx);
+                }
+                // Clear children (already empty for a leaf) so the tombstone
+                // cannot be misidentified as a leaf in the next pass.
+                self.inner.nodes[leaf_idx].children.clear();
+            }
+        }
+    }
+
+    /// Enumerate all root-to-leaf token sequences (the candidate hypotheses
+    /// that will be sent to the target model for verification).
+    #[must_use]
+    pub fn all_paths(&self) -> Vec<Vec<u32>> {
+        self.inner.get_all_paths()
+    }
+
+    /// Mark the longest accepted prefix of `path` and return the number of
+    /// accepted tokens.
+    ///
+    /// The method walks the tree from the root, following `path` one token at
+    /// a time.  Each node whose `token_id` matches the corresponding element
+    /// of `path` is marked `accepted = true`.  When a mismatch is found — or
+    /// when the path is exhausted — we record `actual_next_token` as accepted
+    /// on the mismatch node if one exists and is reachable, then stop.
+    ///
+    /// Returns the count of nodes marked accepted (≥ 0).
+    pub fn accept_path(&mut self, path: &[u32], actual_next_token: u32) -> usize {
+        if path.is_empty() {
+            return 0;
+        }
+
+        let mut accepted = 0usize;
+        // Start at root candidates.
+        let roots: Vec<usize> = self.inner.roots.clone();
+
+        // Find the root node matching path[0].
+        let mut current_idx: Option<usize> = roots
+            .into_iter()
+            .find(|&r| self.inner.nodes[r].token_id == path[0]);
+
+        for (step, &expected_token) in path.iter().enumerate() {
+            match current_idx {
+                None => break,
+                Some(idx) => {
+                    if self.inner.nodes[idx].token_id == expected_token {
+                        self.inner.nodes[idx].accepted = true;
+                        accepted += 1;
+                        // Advance to the matching child for the next step.
+                        if step + 1 < path.len() {
+                            let next_token = path[step + 1];
+                            let children: Vec<usize> = self.inner.nodes[idx].children.clone();
+                            current_idx = children
+                                .into_iter()
+                                .find(|&c| self.inner.nodes[c].token_id == next_token);
+                        } else {
+                            // Path exhausted — try to accept `actual_next_token`
+                            // if a child matches.
+                            let children: Vec<usize> = self.inner.nodes[idx].children.clone();
+                            if let Some(next_idx) = children
+                                .into_iter()
+                                .find(|&c| self.inner.nodes[c].token_id == actual_next_token)
+                            {
+                                self.inner.nodes[next_idx].accepted = true;
+                                accepted += 1;
+                            }
+                            break;
+                        }
+                    } else {
+                        // Mismatch — accept actual_next_token at this position if
+                        // a sibling matches, then stop.
+                        if let Some(parent_idx) = self.inner.nodes[idx].parent {
+                            let siblings: Vec<usize> =
+                                self.inner.nodes[parent_idx].children.clone();
+                            if let Some(sibling_idx) = siblings.into_iter().find(|&s| {
+                                self.inner.nodes[s].token_id == actual_next_token
+                            }) {
+                                self.inner.nodes[sibling_idx].accepted = true;
+                                accepted += 1;
+                            }
+                        }
+                        break;
+                    }
+                }
+            }
+        }
+
+        accepted
+    }
+
+    // ------------------------------------------------------------------
+    // Helpers
+    // ------------------------------------------------------------------
+
+    /// Return up to `k` candidates sorted by probability descending.
+    fn top_k_sorted(candidates: &[(u32, f32)], k: usize) -> Vec<(u32, f32)> {
+        let mut sorted: Vec<(u32, f32)> = candidates.to_vec();
+        // Sort descending by probability; use total_cmp for NaN-safe ordering.
+        sorted.sort_unstable_by(|a, b| b.1.total_cmp(&a.1));
+        sorted.truncate(k);
+        sorted
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Eagle3Decoder
+// ---------------------------------------------------------------------------
+
+/// EAGLE-3 speculative decoder: builds a dynamic draft tree guided by token
+/// confidence.
+///
+/// `Eagle3Decoder` is intentionally stateless — all mutable state lives in
+/// the `DynamicDraftTree` it returns.  This makes it trivially `Clone` and
+/// easy to use from multiple async tasks.
+#[derive(Debug, Clone)]
+pub struct Eagle3Decoder {
+    config: Eagle3Config,
+}
+
+impl Eagle3Decoder {
+    /// Create a new decoder with the given EAGLE-3 configuration.
+    #[must_use]
+    pub fn new(config: Eagle3Config) -> Self {
+        Self { config }
+    }
+
+    /// Build a dynamic draft tree by expanding the draft model iteratively.
+    ///
+    /// # Algorithm
+    ///
+    /// 1. Obtain root candidates by calling
+    ///    `draft_model.draft_from_hidden(hidden_states, context, eagle_config)`.
+    /// 2. For every node whose probability exceeds `expansion_threshold` (and
+    ///    whose depth < `max_depth`), recursively expand using the token
+    ///    appended to `context` as the new context.  Hidden states are
+    ///    approximated by scaling the parent hidden states by the token
+    ///    probability (a cheap proxy; the real implementation would run a
+    ///    lightweight draft head forward pass).
+    /// 3. After the tree reaches `min_depth` everywhere, continue only where
+    ///    `should_expand` returns `true`.
+    /// 4. If `self_consistency` is enabled, call `prune_by_score` with
+    ///    `prune_threshold` to remove low-probability branches before
+    ///    returning.
+    ///
+    /// # Errors
+    ///
+    /// Propagates any error returned by the draft model.
+    pub async fn build_draft_tree(
+        &self,
+        draft_model: &dyn EagleDraftModel,
+        context: &[u32],
+        hidden_states: &[f32],
+    ) -> InferenceResult<DynamicDraftTree> {
+        let mut tree = DynamicDraftTree::new(self.config.clone());
+
+        // Bridge EAGLE-3 config into the legacy EagleConfig the trait expects.
+        let eagle_cfg = EagleConfig {
+            draft_steps: self.config.max_depth,
+            hidden_dim: draft_model.hidden_dim(),
+            num_draft_layers: 1,
+            fusion_method: super::eagle::FusionMethod::Concat,
+            top_k: self.config.beam_width,
+        };
+
+        // --- Depth 0: root candidates ---
+        let root_tokens = draft_model
+            .draft_from_hidden(hidden_states, context, &eagle_cfg)
+            .await?;
+
+        let root_candidates: Vec<(u32, f32)> = root_tokens
+            .iter()
+            .map(|t| (t.id, t.probability))
+            .collect();
+
+        let root_indices = tree.add_roots(&root_candidates);
+
+        // --- Recursive expansion via an explicit work-list (avoids async recursion) ---
+        // Each entry: (node_idx, accumulated_context, approximate_hidden_states)
+        let mut frontier: Vec<(usize, Vec<u32>, Vec<f32>)> = root_indices
+            .into_iter()
+            .map(|idx| {
+                let token_id = tree.inner.nodes[idx].token_id;
+                let mut ctx = context.to_vec();
+                ctx.push(token_id);
+                let prob = tree.inner.nodes[idx].probability;
+                let scaled_hidden: Vec<f32> = hidden_states.iter().map(|h| h * prob).collect();
+                (idx, ctx, scaled_hidden)
+            })
+            .collect();
+
+        // We must always reach min_depth; after that we only expand nodes
+        // that pass `should_expand`.
+        //
+        // Depth semantics: root nodes are at depth 0.  A path of length
+        // `min_depth` has its deepest token at depth `min_depth - 1`.  We
+        // therefore *must* expand a node at depth `d` when `d < min_depth - 1`
+        // (its children will be at `d+1`; that is still shallower than needed).
+        while let Some((parent_idx, ctx, hs)) = {
+            if frontier.is_empty() {
+                None
+            } else {
+                Some(frontier.remove(0))
+            }
+        } {
+            let node_depth = tree.inner.nodes[parent_idx].depth;
+            // Force expansion until each path is at least min_depth tokens.
+            let must_expand = self.config.min_depth > 0
+                && node_depth < self.config.min_depth - 1;
+            let may_expand = tree.should_expand(parent_idx);
+
+            if !must_expand && !may_expand {
+                continue;
+            }
+
+            let child_tokens = draft_model
+                .draft_from_hidden(&hs, &ctx, &eagle_cfg)
+                .await?;
+
+            let child_candidates: Vec<(u32, f32)> = child_tokens
+                .iter()
+                .map(|t| (t.id, t.probability))
+                .collect();
+
+            let child_indices = tree.expand_node(parent_idx, &child_candidates);
+
+            for child_idx in child_indices {
+                let token_id = tree.inner.nodes[child_idx].token_id;
+                let prob = tree.inner.nodes[child_idx].probability;
+                let mut child_ctx = ctx.clone();
+                child_ctx.push(token_id);
+                let child_hs: Vec<f32> = hs.iter().map(|h| h * prob).collect();
+                frontier.push((child_idx, child_ctx, child_hs));
+            }
+        }
+
+        // --- Self-consistency pruning ---
+        if self.config.self_consistency {
+            tree.prune_by_score(self.config.prune_threshold);
+        }
+
+        Ok(tree)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Unit tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::speculative::eagle::{EagleConfig, EagleDraftModel};
+    use crate::speculative::types::Token;
+    use crate::InferenceResult;
+
+    // -----------------------------------------------------------------------
+    // Mock EagleDraftModel
+    // -----------------------------------------------------------------------
+
+    /// A mock draft model that returns a fixed list of `(token_id, probability)`
+    /// pairs on every call.  The same candidates are returned regardless of
+    /// context or hidden states, which makes tree-shape predictable in tests.
+    struct FixedProbDraftModel {
+        /// Candidates returned at every node expansion.
+        candidates: Vec<(u32, f32)>,
+        hidden: usize,
+    }
+
+    impl FixedProbDraftModel {
+        fn new(candidates: Vec<(u32, f32)>, hidden: usize) -> Self {
+            Self { candidates, hidden }
+        }
+    }
+
+    #[async_trait::async_trait]
+    impl EagleDraftModel for FixedProbDraftModel {
+        async fn draft_from_hidden(
+            &self,
+            _hidden_states: &[f32],
+            _context_tokens: &[u32],
+            _config: &EagleConfig,
+        ) -> InferenceResult<Vec<Token>> {
+            Ok(self
+                .candidates
+                .iter()
+                .map(|&(id, prob)| Token {
+                    id,
+                    text: format!("tok_{id}"),
+                    logits: vec![],
+                    probability: prob,
+                })
+                .collect())
+        }
+
+        fn hidden_dim(&self) -> usize {
+            self.hidden
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Helper: build a config that disables self-consistency for structural tests
+    // -----------------------------------------------------------------------
+
+    fn no_prune_config(
+        min_depth: usize,
+        max_depth: usize,
+        threshold: f32,
+        beam: usize,
+    ) -> Eagle3Config {
+        Eagle3Config {
+            min_depth,
+            max_depth,
+            expansion_threshold: threshold,
+            beam_width: beam,
+            self_consistency: false,
+            prune_threshold: 0.0,
+            temperature: 1.0,
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 1 — high-confidence nodes expand to max_depth
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn test_dynamic_tree_expands_high_confidence_nodes() {
+        // All candidates have probability 0.9, well above any threshold.
+        let model = FixedProbDraftModel::new(vec![(1, 0.9)], 4);
+        let config = no_prune_config(1, 4, 0.4, 2);
+        let decoder = Eagle3Decoder::new(config.clone());
+
+        let tree = decoder
+            .build_draft_tree(&model, &[0], &[0.5_f32; 4])
+            .await
+            .unwrap();
+
+        // With beam_width=1 and p=0.9 > 0.4, every node expands until depth 4.
+        // The single path should be depth 4 (4 tokens deep from context).
+        let paths = tree.all_paths();
+        assert!(!paths.is_empty(), "tree must have at least one path");
+        // Every path should reach max_depth tokens.
+        for path in &paths {
+            assert_eq!(
+                path.len(),
+                config.max_depth,
+                "path length should equal max_depth"
+            );
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 2 — low-confidence nodes stay at min_depth
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn test_dynamic_tree_stays_shallow_for_low_confidence() {
+        // Candidates have probability 0.1, below expansion_threshold=0.4.
+        let model = FixedProbDraftModel::new(vec![(1, 0.1)], 4);
+        let config = no_prune_config(1, 6, 0.4, 2);
+        let decoder = Eagle3Decoder::new(config.clone());
+
+        let tree = decoder
+            .build_draft_tree(&model, &[0], &[0.5_f32; 4])
+            .await
+            .unwrap();
+
+        // With p=0.1, `should_expand` is false after min_depth, so all paths
+        // are exactly min_depth long.
+        let paths = tree.all_paths();
+        assert!(!paths.is_empty());
+        for path in &paths {
+            assert_eq!(
+                path.len(),
+                config.min_depth,
+                "paths should stop at min_depth when confidence is low"
+            );
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 3 — prune_by_score removes low-probability branches
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_prune_removes_low_probability_branches() {
+        let config = Eagle3Config {
+            min_depth: 1,
+            max_depth: 3,
+            expansion_threshold: 0.3,
+            beam_width: 3,
+            self_consistency: true,
+            prune_threshold: 0.10,
+            temperature: 1.0,
+        };
+        let mut tree = DynamicDraftTree::new(config);
+
+        // Root: two candidates
+        let r0 = tree.inner.add_root(10, 0.9); // high probability
+        let r1 = tree.inner.add_root(11, 0.05); // below prune_threshold
+
+        // Add one child under r1 to make r1 an internal node rather than a
+        // leaf — prune_by_score only removes leaves, so we test that the
+        // actual low-probability leaf (r1's child) is pruned.
+        let _c0 = tree.inner.add_child(r0, 20, 0.8);
+        let _c1 = tree.inner.add_child(r1, 21, 0.8);
+
+        let before = tree.all_paths().len();
+        tree.prune_by_score(0.10);
+        let after = tree.all_paths().len();
+
+        // The leaf under r1 has path probability 0.05 * 0.8 = 0.04 < 0.10 → pruned.
+        // The leaf under r0 has path probability 0.9 * 0.8 = 0.72 → kept.
+        assert!(
+            after < before,
+            "pruning should have removed at least one path (before={before}, after={after})"
+        );
+        assert_eq!(after, 1, "only the high-probability path should survive");
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 4 — all_paths correct for depth-2, beam-2 tree
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn test_all_paths_correct_for_depth_2_beam_2() {
+        // Two candidates at each level, all with p=0.9 so they expand.
+        let model = FixedProbDraftModel::new(vec![(1, 0.9), (2, 0.85)], 4);
+        let config = no_prune_config(2, 2, 0.4, 2);
+        let decoder = Eagle3Decoder::new(config);
+
+        let tree = decoder
+            .build_draft_tree(&model, &[0], &[1.0_f32; 4])
+            .await
+            .unwrap();
+
+        let paths = tree.all_paths();
+        // 2 roots × 2 children = 4 paths, each of length 2.
+        assert_eq!(paths.len(), 4, "expected 4 paths for beam=2, depth=2");
+        for path in &paths {
+            assert_eq!(path.len(), 2);
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 5 — accept_path returns correct accepted count
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_accept_path_returns_correct_count() {
+        let config = Eagle3Config::default();
+        let mut tree = DynamicDraftTree::new(config);
+
+        let r = tree.inner.add_root(10, 0.9);
+        let c = tree.inner.add_child(r, 20, 0.8);
+        let _gc = tree.inner.add_child(c, 30, 0.7);
+
+        // Accept path [10, 20] — 2 tokens accepted, then actual_next_token=30
+        // which matches the grandchild, so total = 3.
+        let count = tree.accept_path(&[10, 20], 30);
+        assert_eq!(count, 3, "should accept root + child + actual_next_token");
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 6 — path_probability product is correct
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_path_probability_product_correct() {
+        let mut tree = CandidateTree::new();
+        let r = tree.add_root(1, 0.8);
+        let c = tree.add_child(r, 2, 0.5);
+        let gc = tree.add_child(c, 3, 0.25);
+
+        let prob = tree.path_probability(gc);
+        // 0.8 * 0.5 * 0.25 = 0.1
+        assert!(
+            (prob - 0.1_f32).abs() < 1e-6,
+            "expected 0.1, got {prob}"
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 7 — Eagle3Config::default values
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_eagle3_config_default() {
+        let cfg = Eagle3Config::default();
+        assert_eq!(cfg.min_depth, 1);
+        assert_eq!(cfg.max_depth, 6);
+        assert!((cfg.expansion_threshold - 0.4).abs() < 1e-6);
+        assert_eq!(cfg.beam_width, 3);
+        assert!(cfg.self_consistency);
+        assert!((cfg.prune_threshold - 0.05).abs() < 1e-6);
+        assert!((cfg.temperature - 1.0).abs() < 1e-6);
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 8 — leaves() returns only leaf nodes
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_leaves_returns_only_leaf_nodes() {
+        let mut tree = CandidateTree::new();
+        let r = tree.add_root(1, 0.9);
+        let c1 = tree.add_child(r, 2, 0.8);
+        let c2 = tree.add_child(r, 3, 0.7);
+        let _gc = tree.add_child(c1, 4, 0.6); // c1 is now internal
+
+        let leaves = tree.leaves();
+        // c2 (index) and gc (index) are leaves; r and c1 are internal.
+        assert_eq!(leaves.len(), 2, "expected 2 leaves");
+        assert!(!leaves.contains(&r), "root is not a leaf");
+        assert!(!leaves.contains(&c1), "c1 has a child so not a leaf");
+        assert!(leaves.contains(&c2), "c2 is a leaf");
+        let gc_idx = tree.nodes[c1].children[0];
+        assert!(leaves.contains(&gc_idx), "grandchild is a leaf");
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 9 — accept_path with full mismatch returns 0
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_accept_path_mismatch_returns_zero() {
+        let config = Eagle3Config::default();
+        let mut tree = DynamicDraftTree::new(config);
+        let _r = tree.inner.add_root(99, 0.9);
+
+        // path[0] = 1 does not match root token_id 99.
+        let count = tree.accept_path(&[1, 2], 3);
+        assert_eq!(count, 0, "mismatched path should yield 0 accepted tokens");
+    }
+
+    // -----------------------------------------------------------------------
+    // Test 10 — path_probability on out-of-bounds index returns 0.0
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_path_probability_out_of_bounds() {
+        let tree = CandidateTree::new();
+        assert_eq!(tree.path_probability(999), 0.0);
+    }
+}

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

@@ -14,6 +14,7 @@ mod advanced;
 mod config;
 mod decoder;
 mod eagle;
+pub mod eagle3;
 mod error;
 mod lookahead;
 mod medusa;
@@ -27,6 +28,7 @@ pub use advanced::*;
 pub use config::*;
 pub use decoder::*;
 pub use eagle::*;
+pub use eagle3::{DynamicDraftTree, Eagle3Config, Eagle3Decoder};
 pub use error::*;
 pub use lookahead::*;
 pub use medusa::*;

crates/production/rtx-inference/src/speculative/tree.rs

@@ -185,6 +185,46 @@ impl CandidateTree {
             }
         }
     }
+
+    /// Compute the product of probabilities from root to the given leaf.
+    ///
+    /// If `leaf_idx` is out of bounds, returns `0.0`.
+    #[must_use]
+    pub fn path_probability(&self, leaf_idx: usize) -> f32 {
+        if leaf_idx >= self.nodes.len() {
+            return 0.0;
+        }
+        let mut prob = 1.0_f32;
+        let mut current = leaf_idx;
+        loop {
+            prob *= self.nodes[current].probability;
+            match self.nodes[current].parent {
+                Some(parent) => current = parent,
+                None => break,
+            }
+        }
+        prob
+    }
+
+    /// Return the indices of all leaf nodes reachable from roots (nodes with
+    /// no children that are part of the live tree).
+    ///
+    /// Only nodes reachable from `roots` are considered.  Tombstone nodes
+    /// (pruned nodes left in the `nodes` vec but detached from the tree) are
+    /// excluded because they are no longer reachable from any root.
+    #[must_use]
+    pub fn leaves(&self) -> Vec<usize> {
+        let mut result = Vec::new();
+        let mut stack: Vec<usize> = self.roots.clone();
+        while let Some(idx) = stack.pop() {
+            if self.nodes[idx].children.is_empty() {
+                result.push(idx);
+            } else {
+                stack.extend_from_slice(&self.nodes[idx].children);
+            }
+        }
+        result
+    }
 }
 
 impl Default for CandidateTree {

crates/production/rtx-inference/src/speculative/types.rs

@@ -55,6 +55,8 @@ pub enum DraftModelType {
     CachedPredictions(String),
     /// Custom draft model implementation
     Custom(String),
+    /// EAGLE-3 dynamic draft tree with confidence-guided expansion
+    Eagle3,
 }
 
 /// Decision for accepting or rejecting draft tokens