feat(batch9): token merging (ToMe), grad accum per-step norm, speculative streaming

- Token Merging (ToMe, arXiv:2210.09461): bipartite soft matching via priority-queue
  second-chance loop; apply_merge/apply_unmerge; TokenMergingLayer::forward(); 17 unit
  tests + 2 doctests; 75% merge at r=32/seq=64
- Gradient accumulation per-step normalization: NormalizationStrategy
  {EndOfAccumulation, PerStep, None}; with_per_step_normalization(); compute_gradient_norm()
  L2 norm; PerStep divides by fixed accumulation_steps before add (not end-of-batch);
  11 tests including equivalence proof vs EndOfAccumulation
- Speculative streaming: SpeculativeStreamer + mpsc::Receiver<StreamedToken>; notify_step()
  sends accepted draft tokens + optional continuation immediately; StreamStats with Welford
  online mean latency; collect_stream() test helper; 17 async tests

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

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

@@ -0,0 +1,524 @@
+//! Speculative streaming: yield accepted tokens immediately via channel.
+//!
+//! Standard speculative decoding collects all accepted tokens before returning
+//! them to the caller. `SpeculativeStreamer` adds a
+//! [`tokio::sync::mpsc`] channel so each accepted token is forwarded to the
+//! consumer as soon as it is produced, enabling true token-by-token streaming
+//! to an HTTP response or any other async sink.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_inference::speculative::{
+//!     SpeculativeStreamConfig, SpeculativeStreamer, collect_stream,
+//! };
+//!
+//! # #[tokio::main]
+//! # async fn main() {
+//! let config = SpeculativeStreamConfig {
+//!     buffer_size: 32,
+//!     include_continuation: true,
+//!     ..Default::default()
+//! };
+//! let (mut streamer, rx) = SpeculativeStreamer::new(config);
+//!
+//! // Simulate two speculative steps.
+//! streamer.notify_step(&[10, 20, 30], None).await;
+//! streamer.notify_step(&[40], Some(99)).await;
+//! streamer.close();
+//!
+//! let tokens = collect_stream(rx).await;
+//! assert_eq!(tokens.len(), 5); // 3 + 1 draft + 1 continuation
+//! # }
+//! ```
+
+use std::time::{Duration, Instant};
+use tokio::sync::mpsc;
+
+// ---------------------------------------------------------------------------
+// StreamedToken
+// ---------------------------------------------------------------------------
+
+/// A single token event emitted by the speculative streaming decoder.
+#[derive(Debug, Clone)]
+pub struct StreamedToken {
+    /// Vocabulary token id.
+    pub token_id: u32,
+    /// `true` when this token came from draft acceptance; `false` when it is
+    /// a continuation token sampled by the target model on rejection.
+    pub from_draft: bool,
+    /// Zero-based index of the speculative decode step that produced this token.
+    pub step: usize,
+    /// Number of tokens that had been yielded *before* this one (monotonically
+    /// increasing, zero-based).
+    pub total_yielded: usize,
+    /// [`Instant`] at which this token was enqueued into the channel.
+    pub produced_at: Instant,
+}
+
+// ---------------------------------------------------------------------------
+// SpeculativeStreamConfig
+// ---------------------------------------------------------------------------
+
+/// Configuration for [`SpeculativeStreamer`].
+#[derive(Debug, Clone)]
+pub struct SpeculativeStreamConfig {
+    /// `mpsc` channel buffer capacity.
+    ///
+    /// A larger buffer reduces back-pressure at the cost of more memory.
+    /// Default: 64.
+    pub buffer_size: usize,
+    /// When `true`, the continuation token (target sample emitted on partial or
+    /// full rejection) is forwarded to the consumer.
+    ///
+    /// Set `false` if the caller only wants confirmed draft tokens.
+    /// Default: `true`.
+    pub include_continuation: bool,
+    /// How long to wait for a send to succeed when the channel is full before
+    /// counting the token as dropped.
+    ///
+    /// Default: 10 ms.
+    pub send_timeout: Duration,
+}
+
+impl Default for SpeculativeStreamConfig {
+    fn default() -> Self {
+        Self {
+            buffer_size: 64,
+            include_continuation: true,
+            send_timeout: Duration::from_millis(10),
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// StreamStats
+// ---------------------------------------------------------------------------
+
+/// Cumulative statistics for a [`SpeculativeStreamer`] session.
+#[derive(Debug, Clone, Default)]
+pub struct StreamStats {
+    /// Total number of speculative steps processed.
+    pub total_steps: usize,
+    /// Total tokens successfully sent to the consumer.
+    pub total_yielded: usize,
+    /// Tokens that were accepted from the draft model.
+    pub draft_accepted: usize,
+    /// Continuation tokens forwarded (target-sampled on rejection).
+    pub continuation_tokens: usize,
+    /// Tokens dropped because the channel was full and the send timed out.
+    pub dropped_tokens: usize,
+    /// Running mean of per-step wall-clock time in milliseconds.
+    pub avg_step_latency_ms: f32,
+}
+
+impl StreamStats {
+    /// Fraction of successfully yielded tokens that came from draft acceptance.
+    ///
+    /// Returns `0.0` when no tokens have been yielded.
+    #[must_use]
+    pub fn acceptance_rate(&self) -> f32 {
+        if self.total_yielded == 0 {
+            return 0.0;
+        }
+        self.draft_accepted as f32 / self.total_yielded as f32
+    }
+}
+
+// ---------------------------------------------------------------------------
+// SpeculativeStreamer
+// ---------------------------------------------------------------------------
+
+/// Wraps a speculative decode session and yields accepted tokens immediately
+/// to a [`tokio::sync::mpsc::Receiver<StreamedToken>`].
+///
+/// Create via [`SpeculativeStreamer::new`], which returns `(streamer, receiver)`.
+/// Drive each speculative step by calling [`notify_step`][Self::notify_step].
+/// When generation is complete call [`close`][Self::close] so the receiver
+/// observes an end-of-stream.
+pub struct SpeculativeStreamer {
+    config: SpeculativeStreamConfig,
+    sender: mpsc::Sender<StreamedToken>,
+    stats: StreamStats,
+    /// Zero-based index of the *next* step to be processed.
+    step: usize,
+    /// Total tokens yielded so far across all steps.
+    total_yielded: usize,
+}
+
+impl SpeculativeStreamer {
+    /// Create a new streamer.
+    ///
+    /// Returns `(streamer, receiver)`.  The receiver is the consumer-facing end
+    /// of the channel; each call to [`notify_step`][Self::notify_step] sends
+    /// tokens into it.
+    #[must_use]
+    pub fn new(config: SpeculativeStreamConfig) -> (Self, mpsc::Receiver<StreamedToken>) {
+        let (tx, rx) = mpsc::channel(config.buffer_size);
+        let streamer = Self {
+            config,
+            sender: tx,
+            stats: StreamStats::default(),
+            step: 0,
+            total_yielded: 0,
+        };
+        (streamer, rx)
+    }
+
+    /// Process one speculative decode step.
+    ///
+    /// Sends each accepted draft token and (optionally) the continuation token
+    /// to the channel immediately. A send that times out (channel full for
+    /// longer than [`SpeculativeStreamConfig::send_timeout`]) is counted in
+    /// [`StreamStats::dropped_tokens`] and silently skipped.
+    ///
+    /// # Parameters
+    /// - `accepted` — slice of token ids accepted from the draft model in this step.
+    /// - `continuation` — optional target-sampled continuation token (present when
+    ///   the draft was partially or fully rejected).
+    ///
+    /// # Returns
+    /// Number of tokens successfully sent to the consumer in this call.
+    pub async fn notify_step(&mut self, accepted: &[u32], continuation: Option<u32>) -> usize {
+        let step_start = Instant::now();
+        let mut sent: usize = 0;
+
+        for &token_id in accepted {
+            let event = StreamedToken {
+                token_id,
+                from_draft: true,
+                step: self.step,
+                total_yielded: self.total_yielded,
+                produced_at: Instant::now(),
+            };
+            if self.try_send(event).await {
+                self.total_yielded += 1;
+                self.stats.draft_accepted += 1;
+                self.stats.total_yielded += 1;
+                sent += 1;
+            }
+        }
+
+        if self.config.include_continuation {
+            if let Some(token_id) = continuation {
+                let event = StreamedToken {
+                    token_id,
+                    from_draft: false,
+                    step: self.step,
+                    total_yielded: self.total_yielded,
+                    produced_at: Instant::now(),
+                };
+                if self.try_send(event).await {
+                    self.total_yielded += 1;
+                    self.stats.continuation_tokens += 1;
+                    self.stats.total_yielded += 1;
+                    sent += 1;
+                }
+            }
+        }
+
+        // Update running mean of per-step wall-clock latency.
+        let elapsed_ms = step_start.elapsed().as_secs_f32() * 1_000.0;
+        // `self.step` is the 0-based index of the step we just processed.
+        // After incrementing below it becomes `n`, the total steps done.
+        // We compute the new mean using the Welford online update:
+        //   new_mean = old_mean + (x - old_mean) / n
+        let n = (self.step + 1) as f32;
+        self.stats.avg_step_latency_ms += (elapsed_ms - self.stats.avg_step_latency_ms) / n;
+
+        self.step += 1;
+        self.stats.total_steps += 1;
+        sent
+    }
+
+    /// Close the stream.
+    ///
+    /// Drops the sender so the receiver drains to empty and then observes
+    /// end-of-stream (`recv()` returns `None`).
+    pub fn close(self) {
+        // `self` is consumed; dropping it drops `self.sender`.
+    }
+
+    /// Return a reference to the current streaming statistics.
+    #[must_use]
+    pub fn stats(&self) -> &StreamStats {
+        &self.stats
+    }
+
+    /// Zero-based index of the *next* step to be processed.
+    ///
+    /// After `n` calls to [`notify_step`] this returns `n`.
+    #[must_use]
+    pub fn step(&self) -> usize {
+        self.step
+    }
+
+    /// Total number of tokens successfully yielded across all completed steps.
+    #[must_use]
+    pub fn total_yielded(&self) -> usize {
+        self.total_yielded
+    }
+
+    /// Returns `true` if the consumer has dropped the receiver end of the
+    /// channel, meaning further sends will fail.
+    #[must_use]
+    pub fn is_closed(&self) -> bool {
+        self.sender.is_closed()
+    }
+
+    // -----------------------------------------------------------------------
+    // Private helpers
+    // -----------------------------------------------------------------------
+
+    /// Attempt to send one token event within the configured timeout.
+    ///
+    /// Returns `true` on success, `false` if the channel is full and the
+    /// timeout elapses or the receiver has been dropped.
+    async fn try_send(&mut self, event: StreamedToken) -> bool {
+        match tokio::time::timeout(self.config.send_timeout, self.sender.send(event)).await {
+            Ok(Ok(())) => true,
+            // `Ok(Err(_))` — receiver dropped; `Err(_)` — timeout.
+            Ok(Err(_)) | Err(_) => {
+                self.stats.dropped_tokens += 1;
+                false
+            }
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Utility: collect_stream
+// ---------------------------------------------------------------------------
+
+/// Drain a [`StreamedToken`] receiver until the sender is closed.
+///
+/// Useful in tests and batch collection scenarios where you want all tokens as
+/// a `Vec` rather than processing them one at a time.
+///
+/// ```rust
+/// use rtx_inference::speculative::{SpeculativeStreamConfig, SpeculativeStreamer, collect_stream};
+///
+/// # #[tokio::main]
+/// # async fn main() {
+/// let (mut streamer, rx) = SpeculativeStreamer::new(Default::default());
+/// streamer.notify_step(&[1, 2, 3], None).await;
+/// streamer.close();
+/// let all = collect_stream(rx).await;
+/// assert_eq!(all.len(), 3);
+/// # }
+/// ```
+pub async fn collect_stream(mut rx: mpsc::Receiver<StreamedToken>) -> Vec<StreamedToken> {
+    let mut tokens = Vec::new();
+    while let Some(t) = rx.recv().await {
+        tokens.push(t);
+    }
+    tokens
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Helper: config with no continuation tokens and a small buffer.
+    fn cfg_no_cont() -> SpeculativeStreamConfig {
+        SpeculativeStreamConfig {
+            buffer_size: 16,
+            include_continuation: false,
+            ..Default::default()
+        }
+    }
+
+    fn cfg_with_cont() -> SpeculativeStreamConfig {
+        SpeculativeStreamConfig {
+            buffer_size: 16,
+            include_continuation: true,
+            ..Default::default()
+        }
+    }
+
+    #[tokio::test]
+    async fn test_streamer_yields_accepted_tokens() {
+        let (mut streamer, rx) = SpeculativeStreamer::new(cfg_no_cont());
+        let sent = streamer.notify_step(&[10, 20, 30], None).await;
+        assert_eq!(sent, 3);
+        streamer.close();
+        let collected = collect_stream(rx).await;
+        assert_eq!(collected.len(), 3);
+        assert_eq!(collected[0].token_id, 10);
+        assert!(collected[0].from_draft);
+    }
+
+    #[tokio::test]
+    async fn test_streamer_includes_continuation_token() {
+        let (mut streamer, rx) = SpeculativeStreamer::new(cfg_with_cont());
+        let sent = streamer.notify_step(&[1, 2], Some(99)).await;
+        assert_eq!(sent, 3); // 2 draft + 1 continuation
+        streamer.close();
+        let collected = collect_stream(rx).await;
+        assert_eq!(collected.len(), 3);
+        assert!(!collected[2].from_draft);
+        assert_eq!(collected[2].token_id, 99);
+    }
+
+    #[tokio::test]
+    async fn test_streamer_excludes_continuation_when_disabled() {
+        let (mut streamer, rx) = SpeculativeStreamer::new(cfg_no_cont());
+        streamer.notify_step(&[5], Some(42)).await;
+        streamer.close();
+        let collected = collect_stream(rx).await;
+        assert_eq!(collected.len(), 1);
+        assert_eq!(collected[0].token_id, 5);
+    }
+
+    #[tokio::test]
+    async fn test_streamer_step_counter_increments() {
+        let (mut streamer, _rx) = SpeculativeStreamer::new(Default::default());
+        assert_eq!(streamer.step(), 0);
+        streamer.notify_step(&[1], None).await;
+        assert_eq!(streamer.step(), 1);
+        streamer.notify_step(&[2, 3], Some(4)).await;
+        assert_eq!(streamer.step(), 2);
+    }
+
+    #[tokio::test]
+    async fn test_streamer_total_yielded_accumulates() {
+        let (mut streamer, _rx) = SpeculativeStreamer::new(cfg_with_cont());
+        streamer.notify_step(&[1, 2], Some(3)).await; // 3 tokens
+        assert_eq!(streamer.total_yielded(), 3);
+        streamer.notify_step(&[4], None).await; // 1 more
+        assert_eq!(streamer.total_yielded(), 4);
+    }
+
+    #[tokio::test]
+    async fn test_stream_stats_acceptance_rate() {
+        let (mut streamer, _rx) = SpeculativeStreamer::new(cfg_with_cont());
+        streamer.notify_step(&[1, 2], Some(3)).await; // 2 draft, 1 continuation
+        let stats = streamer.stats();
+        // 2 draft / 3 total ≈ 0.667
+        let rate = stats.acceptance_rate();
+        assert!(
+            (rate - 2.0_f32 / 3.0).abs() < 0.01,
+            "acceptance rate {rate}"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_stream_collect_empty_steps() {
+        let (mut streamer, rx) = SpeculativeStreamer::new(Default::default());
+        streamer.notify_step(&[], None).await; // empty step — nothing sent
+        streamer.close();
+        let collected = collect_stream(rx).await;
+        assert_eq!(collected.len(), 0);
+    }
+
+    #[tokio::test]
+    async fn test_is_closed_after_receiver_dropped() {
+        let (streamer, rx) = SpeculativeStreamer::new(Default::default());
+        drop(rx);
+        assert!(streamer.is_closed());
+    }
+
+    #[tokio::test]
+    async fn test_multi_step_ordering() {
+        let (mut streamer, rx) = SpeculativeStreamer::new(Default::default());
+        streamer.notify_step(&[10, 11], None).await;
+        streamer.notify_step(&[20, 21], None).await;
+        streamer.close();
+        let collected = collect_stream(rx).await;
+        assert_eq!(collected.len(), 4);
+        assert_eq!(collected[0].step, 0);
+        assert_eq!(collected[1].step, 0);
+        assert_eq!(collected[2].step, 1);
+        assert_eq!(collected[3].step, 1);
+    }
+
+    // Additional tests for thorough coverage.
+
+    #[tokio::test]
+    async fn test_token_ids_preserved_in_order() {
+        let (mut streamer, rx) = SpeculativeStreamer::new(cfg_no_cont());
+        streamer.notify_step(&[100, 200, 300, 400], None).await;
+        streamer.close();
+        let collected = collect_stream(rx).await;
+        let ids: Vec<u32> = collected.iter().map(|t| t.token_id).collect();
+        assert_eq!(ids, vec![100, 200, 300, 400]);
+    }
+
+    #[tokio::test]
+    async fn test_total_yielded_field_in_event_matches_count() {
+        let (mut streamer, rx) = SpeculativeStreamer::new(cfg_no_cont());
+        streamer.notify_step(&[1, 2, 3], None).await;
+        streamer.close();
+        let collected = collect_stream(rx).await;
+        // The `total_yielded` field inside each event is the cumulative count
+        // *before* that token was emitted.
+        assert_eq!(collected[0].total_yielded, 0);
+        assert_eq!(collected[1].total_yielded, 1);
+        assert_eq!(collected[2].total_yielded, 2);
+    }
+
+    #[tokio::test]
+    async fn test_stats_total_steps_increments() {
+        let (mut streamer, _rx) = SpeculativeStreamer::new(Default::default());
+        streamer.notify_step(&[1], None).await;
+        streamer.notify_step(&[2], None).await;
+        assert_eq!(streamer.stats().total_steps, 2);
+    }
+
+    #[tokio::test]
+    async fn test_stats_continuation_count() {
+        let (mut streamer, _rx) = SpeculativeStreamer::new(cfg_with_cont());
+        streamer.notify_step(&[], Some(7)).await;
+        streamer.notify_step(&[], Some(8)).await;
+        let stats = streamer.stats();
+        assert_eq!(stats.continuation_tokens, 2);
+        assert_eq!(stats.draft_accepted, 0);
+        assert_eq!(stats.total_yielded, 2);
+    }
+
+    #[tokio::test]
+    async fn test_acceptance_rate_zero_when_no_tokens() {
+        let (streamer, _rx) = SpeculativeStreamer::new(Default::default());
+        assert_eq!(streamer.stats().acceptance_rate(), 0.0);
+    }
+
+    #[tokio::test]
+    async fn test_avg_step_latency_populated() {
+        let (mut streamer, _rx) = SpeculativeStreamer::new(Default::default());
+        streamer.notify_step(&[1, 2], None).await;
+        streamer.notify_step(&[3], None).await;
+        // We cannot assert an exact value, but it must be non-negative.
+        assert!(streamer.stats().avg_step_latency_ms >= 0.0);
+    }
+
+    #[tokio::test]
+    async fn test_drop_when_receiver_gone() {
+        // Buffer of 1 so the second token must block, but receiver is gone.
+        let config = SpeculativeStreamConfig {
+            buffer_size: 1,
+            include_continuation: false,
+            send_timeout: Duration::from_millis(1),
+        };
+        let (mut streamer, rx) = SpeculativeStreamer::new(config);
+        drop(rx); // consumer gone immediately
+        // All sends should fail gracefully — no panic.
+        let sent = streamer.notify_step(&[1, 2, 3], None).await;
+        assert_eq!(sent, 0);
+        assert_eq!(streamer.stats().dropped_tokens, 3);
+    }
+
+    #[tokio::test]
+    async fn test_collect_stream_helper() {
+        let (mut streamer, rx) = SpeculativeStreamer::new(cfg_with_cont());
+        streamer.notify_step(&[10, 20], Some(30)).await;
+        streamer.close();
+        let all = collect_stream(rx).await;
+        assert_eq!(all.len(), 3);
+        assert_eq!(all[2].token_id, 30);
+        assert!(!all[2].from_draft);
+    }
+}