feat(batch8): multi-token prediction heads, sparse attention, length bucketing

Multi-Token Prediction heads (rtx-transformers/gpt):
- MtpConfig { num_future_tokens=4, loss_weight=0.3 }; MultiTokenPredictionHead
  with k independent [hidden, vocab] weight matrices; forward() → Vec<Tensor>
- compute_loss(): log-softmax NLL for each offset 1..k; weighted by loss_weight;
  MtpLossResult with per_head_losses + is_valid(); 12 tests

Sparse attention (rtx-transformers/layers):
- SparseAttentionMask: local window (radius w), global tokens (first g attend
  all + all attend them), random long-range (r symmetric positions per token)
- LCG seeded for reproducibility; apply_to_scores() masks to -inf; to_additive_bias()
- SparseAttentionLayer::forward_cpu() skips masked pairs early; numerically-stable
  softmax; sparsity 75% at n=512, 87% at n=1024, 97% at n=4096; 14 tests

Sequence length bucketing (rtx-transformers/training):
- LengthGroupedSampler: Fisher-Yates per-bucket shuffle, token-budget batching,
  overflow bucket for long sequences; padding_efficiency() vs baseline_efficiency()
- pack_into_batch(): greedy first-fit packing; naive_padding_ratio() baseline metric
- Measured 2.2× padding reduction on power-law data (26%→67% efficiency); 18 tests

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

crates/training/rtx-transformers/src/architectures/gpt/mod.rs

@@ -17,6 +17,7 @@ pub mod generation;
 pub mod layer_norm;
 pub mod lm_head;
 pub mod models;
+pub mod multi_token_prediction;
 
 // Re-export main types for convenience
 pub use attention::MultiHeadAttention;
@@ -29,3 +30,4 @@ pub use generation::{
 pub use layer_norm::LayerNorm;
 pub use lm_head::GPTLMHead;
 pub use models::{GPTBlock, GPTLMHeadModel, GPTModel};
+pub use multi_token_prediction::{MtpConfig, MtpLossResult, MultiTokenPredictionHead};

crates/training/rtx-transformers/src/architectures/gpt/multi_token_prediction.rs

@@ -0,0 +1,540 @@
+//! Multi-Token Prediction (MTP) heads.
+//!
+//! Reference: arXiv:2404.19737 — "Better & Faster Large Language Models via Multi-Token Prediction"
+//!
+//! MTP trains `k` independent linear heads, one per future token offset, on top of the same
+//! shared backbone hidden states.  At training time each head `i` predicts token at position
+//! `t + i` from the hidden state at position `t`, contributing a weighted auxiliary loss to the
+//! main language-modelling objective.  At inference time the same heads feed speculative-decoding
+//! draft proposals, yielding a 2–3× throughput gain at near-zero extra FLOPs.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::architectures::gpt::multi_token_prediction::{MtpConfig, MultiTokenPredictionHead};
+//! use rtx_tensor::Device;
+//!
+//! let config = MtpConfig { num_future_tokens: 4, hidden_size: 64, vocab_size: 128, ..Default::default() };
+//! let device = Device::Cpu;
+//! let head = MultiTokenPredictionHead::new(config, &device).unwrap();
+//!
+//! let hidden = rtx_tensor::Tensor::randn(&[1, 8, 64], &device).unwrap();
+//! let logits = head.forward(&hidden).unwrap();
+//! assert_eq!(logits.len(), 4);
+//! ```
+
+use crate::{Result, TransformerError};
+use rtx_tensor::{Device, Tensor};
+
+// ---------------------------------------------------------------------------
+// Configuration
+// ---------------------------------------------------------------------------
+
+/// Configuration for Multi-Token Prediction heads.
+#[derive(Debug, Clone)]
+pub struct MtpConfig {
+    /// Number of additional future tokens to predict.
+    ///
+    /// `k = 4` means heads predict tokens at offsets 1, 2, 3, and 4 from each
+    /// position.  The paper recommends k = 4 for both pre-training efficiency and
+    /// speculative-decoding draft quality.
+    pub num_future_tokens: usize,
+
+    /// Hidden state dimension coming from the backbone.
+    pub hidden_size: usize,
+
+    /// Vocabulary size (output dimension of each head).
+    pub vocab_size: usize,
+
+    /// Scalar weight applied to the summed auxiliary MTP loss before it is added
+    /// to the main loss.  The paper uses 0.3 for pre-training.
+    pub loss_weight: f32,
+
+    /// Whether to include a bias term in each prediction head.
+    pub use_bias: bool,
+}
+
+impl Default for MtpConfig {
+    fn default() -> Self {
+        Self {
+            num_future_tokens: 4,
+            hidden_size: 768,
+            vocab_size: 32000,
+            loss_weight: 0.3,
+            use_bias: false,
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Head
+// ---------------------------------------------------------------------------
+
+/// `k` independent linear projection heads, one per future-token offset.
+///
+/// Each head projects `[batch, seq, hidden_size]` to `[batch, seq, vocab_size]`
+/// via a single `[hidden_size, vocab_size]` weight matrix (and optional bias).
+/// Heads are intentionally independent — they do **not** share weights with the
+/// backbone embedding table — allowing each head to specialise for its offset.
+pub struct MultiTokenPredictionHead {
+    /// Resolved configuration.
+    pub config: MtpConfig,
+
+    /// `num_future_tokens` weight matrices, each `[hidden_size, vocab_size]`.
+    ///
+    /// The outer `Option` is present to mirror the `GPTLMHead` convention and to
+    /// support future weight-tying workflows where a head slot may be replaced by
+    /// a shared tensor.
+    pub heads: Vec<Option<Tensor>>,
+
+    /// Optional bias vectors, each `[vocab_size]`.  `None` when `use_bias` is false.
+    biases: Vec<Option<Tensor>>,
+
+    device: Device,
+}
+
+impl MultiTokenPredictionHead {
+    /// Allocate and initialise `k` prediction heads.
+    ///
+    /// Weights are drawn from `N(0, 0.02²)` following the standard GPT
+    /// initialisation scheme (same as `GPTLMHead`).
+    pub fn new(config: MtpConfig, device: &Device) -> Result<Self> {
+        let init_std = 0.02f32;
+        let mut heads = Vec::with_capacity(config.num_future_tokens);
+        let mut biases = Vec::with_capacity(config.num_future_tokens);
+
+        for _ in 0..config.num_future_tokens {
+            let w = Tensor::randn(&[config.hidden_size, config.vocab_size], device)?
+                .mul_scalar(init_std)
+                .map_err(|e| TransformerError::tensor_op(e.to_string()))?;
+            heads.push(Some(w));
+
+            let b = if config.use_bias {
+                Some(
+                    Tensor::zeros([config.vocab_size], device)
+                        .map_err(|e| TransformerError::tensor_op(e.to_string()))?,
+                )
+            } else {
+                None
+            };
+            biases.push(b);
+        }
+
+        Ok(Self {
+            config,
+            heads,
+            biases,
+            device: device.clone(),
+        })
+    }
+
+    // -----------------------------------------------------------------------
+    // Forward
+    // -----------------------------------------------------------------------
+
+    /// Project `hidden_states` through each head and return `k` logit tensors.
+    ///
+    /// # Arguments
+    ///
+    /// * `hidden_states` — backbone output of shape `[batch, seq, hidden_size]`.
+    ///
+    /// # Returns
+    ///
+    /// A `Vec` of `num_future_tokens` tensors, each of shape
+    /// `[batch, seq, vocab_size]`.  Element `i` corresponds to predictions for
+    /// tokens at offset `i + 1`.
+    pub fn forward(&self, hidden_states: &Tensor) -> Result<Vec<Tensor>> {
+        let mut logits = Vec::with_capacity(self.config.num_future_tokens);
+
+        for (idx, head) in self.heads.iter().enumerate() {
+            let w = head.as_ref().ok_or_else(|| {
+                TransformerError::architecture(format!("MTP head {} weight not initialised", idx))
+            })?;
+
+            let l = hidden_states
+                .matmul(w)
+                .map_err(|e| TransformerError::tensor_op(e.to_string()))?;
+
+            let l = if let Some(bias) = &self.biases[idx] {
+                l.add(bias)
+                    .map_err(|e| TransformerError::tensor_op(e.to_string()))?
+            } else {
+                l
+            };
+
+            logits.push(l);
+        }
+
+        Ok(logits)
+    }
+
+    // -----------------------------------------------------------------------
+    // Loss
+    // -----------------------------------------------------------------------
+
+    /// Compute the weighted MTP auxiliary loss.
+    ///
+    /// For each future offset `i ∈ 1..=k`:
+    ///
+    /// ```text
+    /// loss_i = mean(-log p(labels[t+i] | hidden[t]))  for t in 0..seq_len-i
+    /// ```
+    ///
+    /// The returned scalar is `loss_weight * Σ loss_i`.
+    ///
+    /// This is a CPU reference implementation using log-softmax + NLL, matching
+    /// the approach in `GPTLMHead::compute_loss`.  Negative-one labels (`-1`) are
+    /// silently skipped (ignore index).
+    ///
+    /// # Arguments
+    ///
+    /// * `hidden_states` — `[batch, seq_len, hidden_size]`
+    /// * `labels` — flat `[batch * seq_len]` token ids in row-major order.
+    /// * `batch_size`, `seq_len` — dimensions matching `hidden_states`.
+    pub fn compute_loss(
+        &self,
+        hidden_states: &Tensor,
+        labels: &[i64],
+        batch_size: usize,
+        seq_len: usize,
+    ) -> Result<MtpLossResult> {
+        let all_logits = self.forward(hidden_states)?;
+
+        let vocab = self.config.vocab_size;
+        let mut total_aux_loss = 0.0f32;
+        let mut per_head_losses = Vec::with_capacity(self.config.num_future_tokens);
+
+        for (offset_idx, logits) in all_logits.iter().enumerate() {
+            let offset = offset_idx + 1; // head i predicts token at t + offset
+
+            let valid_len = seq_len.saturating_sub(offset);
+            if valid_len == 0 {
+                per_head_losses.push(0.0f32);
+                continue;
+            }
+
+            // Pull logit data to host.  `to_cpu()` returns Vec<f32>.
+            let logit_data: Vec<f32> = logits.to_cpu().unwrap_or_default();
+
+            let mut head_loss = 0.0f32;
+            let mut count = 0usize;
+
+            for b in 0..batch_size {
+                for t in 0..valid_len {
+                    let label_flat = b * seq_len + t + offset;
+                    if label_flat >= labels.len() {
+                        continue;
+                    }
+                    let label = labels[label_flat];
+                    // Skip ignore-index tokens (convention: -1 or -100)
+                    if label < 0 || label as usize >= vocab {
+                        continue;
+                    }
+
+                    let logit_off = (b * seq_len + t) * vocab;
+                    if logit_off + vocab > logit_data.len() {
+                        continue;
+                    }
+
+                    let row = &logit_data[logit_off..logit_off + vocab];
+
+                    // Numerically stable log-softmax
+                    let max_v = row
+                        .iter()
+                        .copied()
+                        .fold(f32::NEG_INFINITY, f32::max);
+                    let sum_exp: f32 = row.iter().map(|&v| (v - max_v).exp()).sum();
+                    let log_prob = row[label as usize] - max_v - sum_exp.ln();
+
+                    head_loss -= log_prob; // NLL accumulation
+                    count += 1;
+                }
+            }
+
+            let avg_loss = if count > 0 {
+                head_loss / count as f32
+            } else {
+                0.0f32
+            };
+
+            per_head_losses.push(avg_loss);
+            total_aux_loss += avg_loss;
+        }
+
+        let weighted_loss = total_aux_loss * self.config.loss_weight;
+
+        Ok(MtpLossResult {
+            weighted_loss,
+            per_head_losses,
+            num_heads: self.config.num_future_tokens,
+        })
+    }
+
+    // -----------------------------------------------------------------------
+    // Accessors
+    // -----------------------------------------------------------------------
+
+    /// Number of future-token prediction heads (`k`).
+    pub fn num_heads(&self) -> usize {
+        self.config.num_future_tokens
+    }
+
+    /// `(hidden_size, vocab_size)` — the shape of each head's weight matrix.
+    pub fn head_dim(&self) -> (usize, usize) {
+        (self.config.hidden_size, self.config.vocab_size)
+    }
+
+    /// Reference to the device this module lives on.
+    pub fn device(&self) -> &Device {
+        &self.device
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Loss result
+// ---------------------------------------------------------------------------
+
+/// Structured output from [`MultiTokenPredictionHead::compute_loss`].
+#[derive(Debug, Clone)]
+pub struct MtpLossResult {
+    /// `loss_weight * Σ per_head_losses` — add this to the main LM loss.
+    pub weighted_loss: f32,
+
+    /// Unweighted per-head loss, one entry per future offset (1-indexed in
+    /// semantics, 0-indexed in the Vec).
+    pub per_head_losses: Vec<f32>,
+
+    /// Number of heads (mirrors `MtpConfig::num_future_tokens`).
+    pub num_heads: usize,
+}
+
+impl MtpLossResult {
+    /// Returns `true` iff `weighted_loss` is finite (not NaN or ±Inf).
+    pub fn is_valid(&self) -> bool {
+        self.weighted_loss.is_finite()
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rtx_tensor::Device;
+
+    fn cpu() -> Device {
+        Device::Cpu
+    }
+
+    // -- Construction --------------------------------------------------------
+
+    #[test]
+    fn test_mtp_head_creation() {
+        let config = MtpConfig {
+            num_future_tokens: 3,
+            hidden_size: 16,
+            vocab_size: 32,
+            ..Default::default()
+        };
+        let head = MultiTokenPredictionHead::new(config, &cpu()).unwrap();
+        assert_eq!(head.num_heads(), 3);
+        assert_eq!(head.head_dim(), (16, 32));
+    }
+
+    #[test]
+    fn test_mtp_config_default() {
+        let cfg = MtpConfig::default();
+        assert_eq!(cfg.num_future_tokens, 4);
+        assert!((cfg.loss_weight - 0.3).abs() < 1e-6);
+        assert!(!cfg.use_bias);
+    }
+
+    // -- Forward -------------------------------------------------------------
+
+    #[test]
+    fn test_mtp_forward_returns_k_logits() {
+        let config = MtpConfig {
+            num_future_tokens: 4,
+            hidden_size: 8,
+            vocab_size: 16,
+            ..Default::default()
+        };
+        let head = MultiTokenPredictionHead::new(config, &cpu()).unwrap();
+        let hidden = Tensor::randn(&[2, 10, 8], &cpu()).unwrap();
+        let logits = head.forward(&hidden).unwrap();
+        assert_eq!(logits.len(), 4);
+        for l in &logits {
+            let shape = l.shape().dims();
+            // Each head: [batch=2, seq=10, vocab=16]
+            assert_eq!(shape, vec![2, 10, 16]);
+        }
+    }
+
+    #[test]
+    fn test_mtp_heads_are_independent() {
+        // Independent weight matrices must produce distinct outputs.
+        let config = MtpConfig {
+            num_future_tokens: 2,
+            hidden_size: 8,
+            vocab_size: 16,
+            ..Default::default()
+        };
+        let head = MultiTokenPredictionHead::new(config, &cpu()).unwrap();
+        let hidden = Tensor::randn(&[1, 5, 8], &cpu()).unwrap();
+        let logits = head.forward(&hidden).unwrap();
+
+        let l0 = logits[0].to_cpu().unwrap();
+        let l1 = logits[1].to_cpu().unwrap();
+        let different = l0
+            .iter()
+            .zip(l1.iter())
+            .any(|(a, b)| (a - b).abs() > 1e-6);
+        assert!(different, "independent heads must produce different logits");
+    }
+
+    #[test]
+    fn test_mtp_single_head_forward() {
+        let config = MtpConfig {
+            num_future_tokens: 1,
+            hidden_size: 4,
+            vocab_size: 8,
+            ..Default::default()
+        };
+        let head = MultiTokenPredictionHead::new(config, &cpu()).unwrap();
+        let hidden = Tensor::randn(&[1, 3, 4], &cpu()).unwrap();
+        let logits = head.forward(&hidden).unwrap();
+        assert_eq!(logits.len(), 1);
+        assert_eq!(logits[0].shape().dims(), vec![1, 3, 8]);
+    }
+
+    // -- Loss ----------------------------------------------------------------
+
+    #[test]
+    fn test_mtp_loss_is_finite() {
+        let config = MtpConfig {
+            num_future_tokens: 2,
+            hidden_size: 8,
+            vocab_size: 16,
+            loss_weight: 0.3,
+            ..Default::default()
+        };
+        let head = MultiTokenPredictionHead::new(config, &cpu()).unwrap();
+        let hidden = Tensor::randn(&[1, 6, 8], &cpu()).unwrap();
+        let labels: Vec<i64> = (0..6).map(|i| (i % 16) as i64).collect();
+
+        let result = head.compute_loss(&hidden, &labels, 1, 6).unwrap();
+
+        assert!(result.is_valid(), "weighted_loss must be finite");
+        assert_eq!(result.per_head_losses.len(), 2);
+        assert!(result.weighted_loss >= 0.0);
+    }
+
+    #[test]
+    fn test_mtp_loss_weight_scaling() {
+        // weighted_loss == loss_weight * sum(per_head_losses)
+        let config = MtpConfig {
+            num_future_tokens: 1,
+            hidden_size: 8,
+            vocab_size: 16,
+            loss_weight: 0.5,
+            ..Default::default()
+        };
+        let head = MultiTokenPredictionHead::new(config, &cpu()).unwrap();
+        let hidden = Tensor::randn(&[1, 4, 8], &cpu()).unwrap();
+        let labels: Vec<i64> = vec![0, 1, 2, 3];
+
+        let result = head.compute_loss(&hidden, &labels, 1, 4).unwrap();
+        let expected = result.per_head_losses[0] * 0.5;
+        assert!(
+            (result.weighted_loss - expected).abs() < 1e-5,
+            "weighted_loss={} expected={}",
+            result.weighted_loss,
+            expected
+        );
+    }
+
+    #[test]
+    fn test_mtp_zero_offset_skips_gracefully() {
+        // With seq_len=2 and k=4, heads 3 and 4 (offsets 3, 4) have valid_len=0
+        // and must contribute 0.0 to avoid corrupting the sum.
+        let config = MtpConfig {
+            num_future_tokens: 4,
+            hidden_size: 8,
+            vocab_size: 16,
+            ..Default::default()
+        };
+        let head = MultiTokenPredictionHead::new(config, &cpu()).unwrap();
+        let hidden = Tensor::randn(&[1, 2, 8], &cpu()).unwrap();
+        let labels: Vec<i64> = vec![1, 2];
+
+        let result = head.compute_loss(&hidden, &labels, 1, 2).unwrap();
+
+        // Offsets 3 (idx 2) and 4 (idx 3) must be zero.
+        assert!(
+            result.per_head_losses[2].abs() < 1e-6,
+            "head 3 loss should be 0, got {}",
+            result.per_head_losses[2]
+        );
+        assert!(
+            result.per_head_losses[3].abs() < 1e-6,
+            "head 4 loss should be 0, got {}",
+            result.per_head_losses[3]
+        );
+    }
+
+    #[test]
+    fn test_mtp_loss_nonnegative() {
+        // Cross-entropy is always >= 0.
+        let config = MtpConfig {
+            num_future_tokens: 3,
+            hidden_size: 8,
+            vocab_size: 16,
+            loss_weight: 0.3,
+            ..Default::default()
+        };
+        let head = MultiTokenPredictionHead::new(config, &cpu()).unwrap();
+        let hidden = Tensor::randn(&[2, 8, 8], &cpu()).unwrap();
+        let labels: Vec<i64> = (0..16).map(|i| (i % 16) as i64).collect();
+
+        let result = head.compute_loss(&hidden, &labels, 2, 8).unwrap();
+
+        for (i, &l) in result.per_head_losses.iter().enumerate() {
+            assert!(l >= 0.0, "head {} loss {} should be non-negative", i, l);
+        }
+        assert!(result.weighted_loss >= 0.0);
+    }
+
+    // -- MtpLossResult -------------------------------------------------------
+
+    #[test]
+    fn test_mtp_loss_result_valid() {
+        let r = MtpLossResult {
+            weighted_loss: 1.23,
+            per_head_losses: vec![1.0, 0.23],
+            num_heads: 2,
+        };
+        assert!(r.is_valid());
+    }
+
+    #[test]
+    fn test_mtp_loss_result_invalid_nan() {
+        let r = MtpLossResult {
+            weighted_loss: f32::NAN,
+            per_head_losses: vec![],
+            num_heads: 0,
+        };
+        assert!(!r.is_valid());
+    }
+
+    #[test]
+    fn test_mtp_loss_result_invalid_inf() {
+        let r = MtpLossResult {
+            weighted_loss: f32::INFINITY,
+            per_head_losses: vec![],
+            num_heads: 0,
+        };
+        assert!(!r.is_valid());
+    }
+}

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

@@ -32,8 +32,8 @@ pub mod llama_attention;
 pub use bert::{BertConfig, BertEmbeddings, BertModel, BertPooler};
 pub use gpt::{
     FeedForward, GPTBlock, GPTConfig, GPTLMHeadModel, GPTModel, GenerationConfig, LayerNorm,
-    MultiHeadAttention, PositionEncodingType, PositionalEmbedding, SamplingStrategy, TextGenerator,
+    MtpConfig, MtpLossResult, MultiHeadAttention, MultiTokenPredictionHead, PositionEncodingType,
-    TokenEmbedding,
+    PositionalEmbedding, SamplingStrategy, TextGenerator, TokenEmbedding,
 };
 pub use llama::{LLaMAConfig, RotaryPositionEmbedding, SwiGLU};
 pub use transformer_block::TransformerBlock;

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

@@ -41,6 +41,7 @@ pub mod glu_activations;
 pub mod flex_attention;
 pub mod multi_query_attention;
 pub mod sage_attention;
+pub mod sparse_attention;
 // pub mod grouped_query_attention;
 // pub mod sliding_window_attention;
 
@@ -224,6 +225,7 @@ pub use ring_attention::{
 pub use sage_attention::{
     QuantParams, SageAttention, SageAttentionBuilder, SageAttentionConfig, SageAttentionOutput,
 };
+pub use sparse_attention::{SparseAttentionConfig, SparseAttentionLayer, SparseAttentionMask};
 
 // TransformerConfig is defined above and available for import
 

crates/training/rtx-transformers/src/layers/sparse_attention.rs

@@ -0,0 +1,653 @@
+//! Sparse attention: local window + global tokens + random long-range positions.
+//!
+//! Implements the BigBird/Longformer attention pattern described in:
+//! - BigBird: arXiv:2007.14062
+//! - Longformer: arXiv:2004.05150
+//!
+//! The full O(n²) attention matrix is replaced by three complementary patterns:
+//!
+//! 1. **Local window** — each token i attends to positions \[i-w, i+w\] (clipped at boundaries).
+//! 2. **Global tokens** — the first `g` tokens attend to every position, and every position
+//!    attends back to them.  These act as "hub" nodes that propagate global context.
+//! 3. **Random long-range** — each token draws `r` additional positions uniformly at random,
+//!    providing O(log n) connectivity for long-range information flow.
+//!
+//! Total attended pairs: O(n·w + g·n + r·n) instead of O(n²).
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::layers::sparse_attention::{SparseAttentionConfig, SparseAttentionLayer};
+//!
+//! let config = SparseAttentionConfig {
+//!     window_radius: 64,
+//!     num_global_tokens: 2,
+//!     num_random_tokens: 3,
+//!     seq_len: 512,
+//! };
+//! let layer = SparseAttentionLayer::new(config, 64 /* head_dim */);
+//! println!("sparsity: {:.2}%", layer.sparsity() * 100.0);
+//! ```
+
+// ---------------------------------------------------------------------------
+// Configuration
+// ---------------------------------------------------------------------------
+
+/// Configuration for a sparse attention mask.
+///
+/// All fields are sequence-length-independent except `seq_len`, which fixes
+/// the concrete mask size.  Create a new mask whenever the sequence length
+/// changes (typically once per batch if lengths are uniform).
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct SparseAttentionConfig {
+    /// One-sided local window radius.
+    ///
+    /// Token i attends to positions in `[i.saturating_sub(w), (i+w+1).min(n))`.
+    /// A radius of 64 means each token sees up to 129 neighbours.
+    pub window_radius: usize,
+
+    /// Number of global tokens at the **start** of the sequence (0-indexed).
+    ///
+    /// Each global token attends to every position, and every position attends
+    /// back to it.  Typical values: 1 (CLS token) or a small constant.
+    pub num_global_tokens: usize,
+
+    /// Extra long-range positions sampled uniformly per token.
+    ///
+    /// The mask is kept **symmetric**: if i draws j then j→i is also set.
+    /// Positions are deterministic given the seed passed to `new_with_seed`.
+    pub num_random_tokens: usize,
+
+    /// Sequence length this mask is built for.
+    pub seq_len: usize,
+}
+
+impl Default for SparseAttentionConfig {
+    fn default() -> Self {
+        Self {
+            window_radius: 64,
+            num_global_tokens: 1,
+            num_random_tokens: 3,
+            seq_len: 512,
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// SparseAttentionMask
+// ---------------------------------------------------------------------------
+
+/// Dense boolean representation of a sparse attention pattern.
+///
+/// `mask[i][j] == true` means token i is allowed to attend to token j.
+///
+/// For large sequences this `Vec<Vec<bool>>` is a CPU-side reference
+/// implementation.  A production path would materialise only the non-zero
+/// indices (CSR/COO) and feed them to a sparse CUDA/Metal kernel.
+pub struct SparseAttentionMask {
+    /// The configuration used to build this mask.
+    pub config: SparseAttentionConfig,
+    /// Row-major boolean matrix of shape `[seq_len, seq_len]`.
+    pub mask: Vec<Vec<bool>>,
+}
+
+impl SparseAttentionMask {
+    /// Build the sparse mask seeded with `42` for reproducibility.
+    ///
+    /// Equivalent to `Self::new_with_seed(config, 42)`.
+    pub fn new(config: SparseAttentionConfig) -> Self {
+        Self::new_with_seed(config, 42)
+    }
+
+    /// Build the sparse mask with an explicit RNG seed.
+    ///
+    /// The random component uses a Knuth MMIX LCG stepped once per (row, draw)
+    /// pair, so results are fully deterministic and independent of platform.
+    ///
+    /// # Panics
+    ///
+    /// Does not panic; all arithmetic is wrapping.
+    pub fn new_with_seed(config: SparseAttentionConfig, seed: u64) -> Self {
+        let n = config.seq_len;
+        let mut mask = vec![vec![false; n]; n];
+
+        // ----------------------------------------------------------------
+        // 1. Local window: token i attends to [max(0,i-w) .. min(n,i+w+1))
+        // ----------------------------------------------------------------
+        for i in 0..n {
+            let lo = i.saturating_sub(config.window_radius);
+            let hi = (i + config.window_radius + 1).min(n);
+            for j in lo..hi {
+                mask[i][j] = true;
+            }
+        }
+
+        // ----------------------------------------------------------------
+        // 2. Global tokens: first g tokens attend everywhere; all attend to them
+        // ----------------------------------------------------------------
+        for g in 0..config.num_global_tokens.min(n) {
+            for j in 0..n {
+                mask[g][j] = true;
+            }
+            for i in 0..n {
+                mask[i][g] = true;
+            }
+        }
+
+        // ----------------------------------------------------------------
+        // 3. Random long-range positions (symmetric)
+        //
+        // LCG: state = state * 6364136223846793005 + 1442695040888963407
+        // These constants are the Knuth MMIX multiplier/increment.
+        // We advance once per (row i, draw k) pair so the sequence for
+        // each token is independent of window_radius and global_tokens.
+        // ----------------------------------------------------------------
+        let mut rng_state = seed;
+        for i in 0..n {
+            for _ in 0..config.num_random_tokens {
+                rng_state = rng_state
+                    .wrapping_mul(6_364_136_223_846_793_005)
+                    .wrapping_add(1_442_695_040_888_963_407);
+                // Use high bits — they have better statistical properties in LCGs.
+                let j = (rng_state >> 33) as usize % n;
+                mask[i][j] = true;
+                mask[j][i] = true; // enforce symmetry
+            }
+        }
+
+        Self { config, mask }
+    }
+
+    // -----------------------------------------------------------------------
+    // Inspection
+    // -----------------------------------------------------------------------
+
+    /// Number of positions token `i` attends to.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `i >= seq_len`.
+    pub fn attention_count(&self, i: usize) -> usize {
+        self.mask[i].iter().filter(|&&v| v).count()
+    }
+
+    /// Fraction of entries in the attention matrix that are **blocked** (not attended).
+    ///
+    /// Returns a value in `[0.0, 1.0)`.  Higher values mean sparser attention.
+    pub fn sparsity(&self) -> f32 {
+        let n = self.config.seq_len;
+        if n == 0 {
+            return 0.0;
+        }
+        let attended: usize = self.mask.iter().flatten().filter(|&&v| v).count();
+        1.0 - attended as f32 / (n * n) as f32
+    }
+
+    // -----------------------------------------------------------------------
+    // Score modification helpers
+    // -----------------------------------------------------------------------
+
+    /// Return a flat additive bias vector of length `seq_len²`.
+    ///
+    /// Attended positions receive `0.0`; blocked positions receive `-1e9`,
+    /// which after softmax reduces to ~0 probability.
+    pub fn to_additive_bias(&self) -> Vec<f32> {
+        self.mask
+            .iter()
+            .flatten()
+            .map(|&v| if v { 0.0_f32 } else { -1.0e9_f32 })
+            .collect()
+    }
+
+    /// Apply the mask to a flat attention-score slice **in place**.
+    ///
+    /// `scores` must be a row-major `[seq_len, seq_len]` slice.  Blocked
+    /// positions are set to `f32::NEG_INFINITY` so that softmax zeroes them.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `scores.len() != seq_len²`.
+    pub fn apply_to_scores(&self, scores: &mut [f32]) {
+        let n = self.config.seq_len;
+        assert_eq!(
+            scores.len(),
+            n * n,
+            "scores length {} does not match seq_len² = {}",
+            scores.len(),
+            n * n
+        );
+        for i in 0..n {
+            for j in 0..n {
+                if !self.mask[i][j] {
+                    scores[i * n + j] = f32::NEG_INFINITY;
+                }
+            }
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// SparseAttentionLayer
+// ---------------------------------------------------------------------------
+
+/// CPU reference implementation of masked multi-head attention with sparse patterns.
+///
+/// This is a single-head-per-call, pure-Rust forward pass suitable for
+/// correctness testing, ablation studies, and small-sequence inference.
+/// GPU-accelerated paths would consume the same `SparseAttentionMask` but
+/// materialise it as a COO/CSR sparse structure fed into a tiled kernel.
+///
+/// # Layout
+///
+/// Q, K, V tensors are passed as flat `f32` slices with the logical layout
+/// `[num_heads, seq_len, head_dim]`, i.e. head-major, token-major, dimension-minor.
+pub struct SparseAttentionLayer {
+    /// Configuration shared with the embedded mask.
+    pub config: SparseAttentionConfig,
+    /// Pre-built sparse mask for this `(seq_len, window_radius, …)` configuration.
+    pub mask: SparseAttentionMask,
+    /// Scaling factor applied before softmax: `1 / sqrt(head_dim)`.
+    pub scale: f32,
+}
+
+impl SparseAttentionLayer {
+    /// Create a layer from config and head dimension.
+    ///
+    /// The mask is built once at construction and reused across forward calls.
+    pub fn new(config: SparseAttentionConfig, head_dim: usize) -> Self {
+        let scale = 1.0 / (head_dim as f32).sqrt();
+        let mask = SparseAttentionMask::new(config.clone());
+        Self { config, mask, scale }
+    }
+
+    /// CPU reference forward pass.
+    ///
+    /// # Arguments
+    ///
+    /// - `q`, `k`, `v` — flat `f32` slices shaped `[num_heads, seq_len, head_dim]`.
+    /// - `num_heads` — number of attention heads.
+    /// - `head_dim` — dimension of each head.
+    ///
+    /// # Returns
+    ///
+    /// Output tensor shaped `[num_heads, seq_len, head_dim]` (flat `Vec<f32>`).
+    ///
+    /// # Panics
+    ///
+    /// Panics if slice lengths are inconsistent with `num_heads * seq_len * head_dim`.
+    pub fn forward_cpu(
+        &self,
+        q: &[f32],
+        k: &[f32],
+        v: &[f32],
+        num_heads: usize,
+        head_dim: usize,
+    ) -> Vec<f32> {
+        let n = self.config.seq_len;
+        let head_stride = n * head_dim; // elements per head in q/k/v
+
+        assert_eq!(q.len(), num_heads * head_stride);
+        assert_eq!(k.len(), num_heads * head_stride);
+        assert_eq!(v.len(), num_heads * head_stride);
+
+        let mut output = vec![0.0_f32; num_heads * head_stride];
+
+        for h in 0..num_heads {
+            let q_h = &q[h * head_stride..(h + 1) * head_stride];
+            let k_h = &k[h * head_stride..(h + 1) * head_stride];
+            let v_h = &v[h * head_stride..(h + 1) * head_stride];
+            let out_h = &mut output[h * head_stride..(h + 1) * head_stride];
+
+            // ---- QK^T scaled dot product [n, n] --------------------------------
+            let mut scores = vec![0.0_f32; n * n];
+            for i in 0..n {
+                let q_row = &q_h[i * head_dim..(i + 1) * head_dim];
+                for j in 0..n {
+                    // Skip computation for masked positions to save work.
+                    if !self.mask.mask[i][j] {
+                        scores[i * n + j] = f32::NEG_INFINITY;
+                        continue;
+                    }
+                    let k_row = &k_h[j * head_dim..(j + 1) * head_dim];
+                    let dot: f32 = q_row.iter().zip(k_row).map(|(&qi, &ki)| qi * ki).sum();
+                    scores[i * n + j] = dot * self.scale;
+                }
+            }
+
+            // ---- Row-wise numerically-stable softmax ---------------------------
+            for i in 0..n {
+                let row = &mut scores[i * n..(i + 1) * n];
+                // max over attended positions only (NEG_INFINITY entries skipped)
+                let max_v = row
+                    .iter()
+                    .cloned()
+                    .filter(|s| s.is_finite())
+                    .fold(f32::NEG_INFINITY, f32::max);
+
+                if max_v.is_infinite() {
+                    // All positions masked: row stays as zeros (no contribution).
+                    row.iter_mut().for_each(|s| *s = 0.0);
+                    continue;
+                }
+
+                let mut sum = 0.0_f32;
+                for s in row.iter_mut() {
+                    *s = (*s - max_v).exp(); // NEG_INFINITY → exp(-inf) = 0
+                    sum += *s;
+                }
+                if sum > 0.0 {
+                    row.iter_mut().for_each(|s| *s /= sum);
+                }
+            }
+
+            // ---- Weighted sum of V [n, head_dim] --------------------------------
+            for i in 0..n {
+                for d in 0..head_dim {
+                    let mut acc = 0.0_f32;
+                    for j in 0..n {
+                        acc += scores[i * n + j] * v_h[j * head_dim + d];
+                    }
+                    out_h[i * head_dim + d] = acc;
+                }
+            }
+        }
+
+        output
+    }
+
+    /// Sparsity of the underlying attention pattern (fraction of blocked pairs).
+    pub fn sparsity(&self) -> f32 {
+        self.mask.sparsity()
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // -----------------------------------------------------------------------
+    // Mask construction
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_mask_local_window_attended() {
+        let config = SparseAttentionConfig {
+            window_radius: 2,
+            num_global_tokens: 0,
+            num_random_tokens: 0,
+            seq_len: 8,
+        };
+        let mask = SparseAttentionMask::new_with_seed(config, 0);
+        // Token 4 attends to [2,3,4,5,6].
+        assert!(mask.mask[4][2]);
+        assert!(mask.mask[4][3]);
+        assert!(mask.mask[4][4]);
+        assert!(mask.mask[4][5]);
+        assert!(mask.mask[4][6]);
+        // Token 4 must NOT attend to tokens 0 or 1 (outside window, no globals).
+        assert!(!mask.mask[4][0]);
+        assert!(!mask.mask[4][1]);
+    }
+
+    #[test]
+    fn test_mask_global_tokens_attend_all() {
+        let config = SparseAttentionConfig {
+            window_radius: 0,
+            num_global_tokens: 2,
+            num_random_tokens: 0,
+            seq_len: 8,
+        };
+        let mask = SparseAttentionMask::new_with_seed(config, 0);
+        // Global token 0: attends everywhere.
+        for j in 0..8 {
+            assert!(mask.mask[0][j], "global[0] should attend to {j}");
+        }
+        // All tokens attend to global 0.
+        for i in 0..8 {
+            assert!(mask.mask[i][0], "token {i} should attend to global[0]");
+        }
+        // Global token 1 as well.
+        for j in 0..8 {
+            assert!(mask.mask[1][j], "global[1] should attend to {j}");
+        }
+        for i in 0..8 {
+            assert!(mask.mask[i][1], "token {i} should attend to global[1]");
+        }
+    }
+
+    #[test]
+    fn test_mask_sparsity_increases_with_length() {
+        let cfg_short = SparseAttentionConfig {
+            window_radius: 4,
+            num_global_tokens: 1,
+            num_random_tokens: 0,
+            seq_len: 16,
+        };
+        let cfg_long = SparseAttentionConfig {
+            window_radius: 4,
+            num_global_tokens: 1,
+            num_random_tokens: 0,
+            seq_len: 64,
+        };
+        let s_short = SparseAttentionMask::new_with_seed(cfg_short, 0).sparsity();
+        let s_long = SparseAttentionMask::new_with_seed(cfg_long, 0).sparsity();
+        assert!(
+            s_long > s_short,
+            "longer sequence should yield higher sparsity for the same window: {s_long} vs {s_short}"
+        );
+    }
+
+    #[test]
+    fn test_mask_random_tokens_are_symmetric() {
+        let config = SparseAttentionConfig {
+            window_radius: 0,
+            num_global_tokens: 0,
+            num_random_tokens: 2,
+            seq_len: 8,
+        };
+        let mask = SparseAttentionMask::new_with_seed(config, 123);
+        for i in 0..8 {
+            for j in 0..8 {
+                if mask.mask[i][j] {
+                    assert!(
+                        mask.mask[j][i],
+                        "mask[{i}][{j}] is true but mask[{j}][{i}] is false — symmetry broken"
+                    );
+                }
+            }
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Score helpers
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_additive_bias_values() {
+        let config = SparseAttentionConfig {
+            window_radius: 1,
+            num_global_tokens: 0,
+            num_random_tokens: 0,
+            seq_len: 4,
+        };
+        let mask = SparseAttentionMask::new_with_seed(config, 0);
+        let bias = mask.to_additive_bias();
+        // Diagonal (self-attend is inside window) must be 0.
+        assert_eq!(bias[0 * 4 + 0], 0.0, "bias[0][0] should be 0");
+        assert_eq!(bias[1 * 4 + 1], 0.0, "bias[1][1] should be 0");
+        // mask[0][3]: distance 3 > window 1 → blocked.
+        assert!(
+            bias[0 * 4 + 3] < -1.0e8,
+            "bias[0][3] should be large negative, got {}",
+            bias[0 * 4 + 3]
+        );
+    }
+
+    #[test]
+    fn test_apply_to_scores_sets_neg_inf() {
+        let config = SparseAttentionConfig {
+            window_radius: 1,
+            num_global_tokens: 0,
+            num_random_tokens: 0,
+            seq_len: 4,
+        };
+        let mask = SparseAttentionMask::new_with_seed(config, 0);
+        let mut scores = vec![1.0_f32; 16];
+        mask.apply_to_scores(&mut scores);
+        for i in 0..4_usize {
+            for j in 0..4_usize {
+                if !mask.mask[i][j] {
+                    assert!(
+                        scores[i * 4 + j].is_infinite() && scores[i * 4 + j] < 0.0,
+                        "scores[{i}][{j}] should be -inf, got {}",
+                        scores[i * 4 + j]
+                    );
+                }
+            }
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Forward pass
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_sparse_attention_forward_shape() {
+        let config = SparseAttentionConfig {
+            window_radius: 2,
+            num_global_tokens: 1,
+            num_random_tokens: 1,
+            seq_len: 8,
+        };
+        let layer = SparseAttentionLayer::new(config, 4);
+        let num_heads = 2;
+        let head_dim = 4;
+        let seq = 8;
+        let q = vec![0.1_f32; num_heads * seq * head_dim];
+        let k = vec![0.1_f32; num_heads * seq * head_dim];
+        let v = vec![0.5_f32; num_heads * seq * head_dim];
+        let out = layer.forward_cpu(&q, &k, &v, num_heads, head_dim);
+        assert_eq!(out.len(), num_heads * seq * head_dim);
+    }
+
+    #[test]
+    fn test_sparse_attention_global_token_output_not_zero() {
+        // Global token 0 attends to all 4 positions; its output must aggregate V.
+        let config = SparseAttentionConfig {
+            window_radius: 0,
+            num_global_tokens: 1,
+            num_random_tokens: 0,
+            seq_len: 4,
+        };
+        let layer = SparseAttentionLayer::new(config, 4);
+        let q = vec![1.0_f32; 4 * 4]; // 1 head, 4 tokens, dim 4
+        let k = vec![1.0_f32; 4 * 4];
+        let v: Vec<f32> = (0_u32..16).map(|i| i as f32).collect();
+        let out = layer.forward_cpu(&q, &k, &v, 1, 4);
+        let global_out = &out[..4]; // first token output
+        assert!(
+            global_out.iter().any(|&val| val > 0.0),
+            "global token output should be non-zero, got {global_out:?}"
+        );
+    }
+
+    #[test]
+    fn test_sparse_attention_sparsity_positive() {
+        let config = SparseAttentionConfig {
+            window_radius: 2,
+            num_global_tokens: 1,
+            num_random_tokens: 1,
+            seq_len: 32,
+        };
+        let layer = SparseAttentionLayer::new(config, 8);
+        let sp = layer.sparsity();
+        assert!(sp > 0.0, "sparsity should be positive");
+        assert!(sp < 1.0, "sparsity should be below 1.0");
+    }
+
+    #[test]
+    fn test_attention_count_within_bounds() {
+        let config = SparseAttentionConfig {
+            window_radius: 3,
+            num_global_tokens: 2,
+            num_random_tokens: 2,
+            seq_len: 16,
+        };
+        let mask = SparseAttentionMask::new(config.clone());
+        for i in 0..16 {
+            let cnt = mask.attention_count(i);
+            assert!(cnt > 0, "token {i} must attend to at least one position");
+            assert!(cnt <= 16, "token {i} cannot attend to more than seq_len positions");
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Extra invariant checks
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_self_attend_always_set_with_window() {
+        // Window radius ≥ 0 always includes i==j (the diagonal).
+        let config = SparseAttentionConfig {
+            window_radius: 0,
+            num_global_tokens: 0,
+            num_random_tokens: 0,
+            seq_len: 10,
+        };
+        let mask = SparseAttentionMask::new_with_seed(config, 0);
+        for i in 0..10 {
+            assert!(mask.mask[i][i], "self-attention must always be set (window_radius=0)");
+        }
+    }
+
+    #[test]
+    fn test_sparsity_zero_for_dense_window() {
+        // Window radius >= seq_len-1 → full attention → sparsity 0.
+        let n = 8_usize;
+        let config = SparseAttentionConfig {
+            window_radius: n - 1,
+            num_global_tokens: 0,
+            num_random_tokens: 0,
+            seq_len: n,
+        };
+        let mask = SparseAttentionMask::new_with_seed(config, 0);
+        assert_eq!(mask.sparsity(), 0.0, "full window should yield sparsity=0");
+    }
+
+    #[test]
+    fn test_num_global_tokens_clamped_to_seq_len() {
+        // num_global_tokens > seq_len should not panic.
+        let config = SparseAttentionConfig {
+            window_radius: 0,
+            num_global_tokens: 100, // >> seq_len=4
+            num_random_tokens: 0,
+            seq_len: 4,
+        };
+        let mask = SparseAttentionMask::new_with_seed(config, 0);
+        // With g >= n every pair is attended → full matrix.
+        assert_eq!(mask.sparsity(), 0.0);
+    }
+
+    #[test]
+    fn test_forward_output_finite() {
+        // Verify no NaN/Inf leaks into output (important with -inf masking + exp).
+        let config = SparseAttentionConfig {
+            window_radius: 1,
+            num_global_tokens: 1,
+            num_random_tokens: 1,
+            seq_len: 6,
+        };
+        let layer = SparseAttentionLayer::new(config, 4);
+        let q: Vec<f32> = (0..6 * 4).map(|i| (i as f32) * 0.01).collect();
+        let k: Vec<f32> = q.iter().map(|&x| x * 1.1).collect();
+        let v: Vec<f32> = (0..6 * 4).map(|i| (i as f32) * 0.005).collect();
+        let out = layer.forward_cpu(&q, &k, &v, 1, 4);
+        for (idx, &val) in out.iter().enumerate() {
+            assert!(val.is_finite(), "output[{idx}] = {val} is not finite");
+        }
+    }
+}

crates/training/rtx-transformers/src/training/length_bucketing.rs

@@ -0,0 +1,636 @@
+//! Sequence length bucketing for reduced padding waste.
+//!
+//! Groups training samples by sequence length into buckets.
+//! Within each bucket, all sequences are close in length, so padding waste is minimal.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::training::length_bucketing::{
+//!     BucketConfig, LengthGroupedSampler, pack_into_batch,
+//! };
+//!
+//! // Power-law-like dataset: many short sequences, few long ones
+//! let lengths: Vec<usize> = (0..200)
+//!     .map(|i| if i < 150 { 50 + i % 30 } else { 500 + i * 5 })
+//!     .collect();
+//!
+//! let config = BucketConfig {
+//!     shuffle_within_bucket: false,
+//!     ..Default::default()
+//! };
+//! let sampler = LengthGroupedSampler::new(lengths, config);
+//!
+//! let bucketed_eff = sampler.padding_efficiency();
+//! let baseline_eff = sampler.baseline_efficiency();
+//! assert!(bucketed_eff > baseline_eff);
+//! ```
+
+/// Bucketing configuration.
+#[derive(Debug, Clone)]
+pub struct BucketConfig {
+    /// Bucket boundary lengths.
+    ///
+    /// Sequences with length <= `boundary[i]` are placed in bucket `i`.
+    /// The final bucket (index `boundaries.len()`) catches all longer sequences.
+    ///
+    /// Default: `[64, 128, 256, 512, 1024, 2048]`.
+    pub bucket_boundaries: Vec<usize>,
+    /// Maximum total tokens per batch (dynamic batch size).
+    ///
+    /// The number of sequences per batch is `max_tokens_per_batch / padded_bucket_length`.
+    pub max_tokens_per_batch: usize,
+    /// Pad all sequences in a batch to the nearest multiple of this value.
+    ///
+    /// Values <= 1 disable alignment. Hardware-friendly choices: 8 (FP16 Tensor Cores),
+    /// 64 (cache lines).
+    pub pad_to_multiple_of: usize,
+    /// Shuffle sequences within each bucket (for training).
+    ///
+    /// Uses a deterministic LCG seeded by [`BucketConfig::seed`]; disable for
+    /// reproducible ordering in evaluation.
+    pub shuffle_within_bucket: bool,
+    /// Fixed seed for reproducible shuffling.
+    pub seed: u64,
+}
+
+impl Default for BucketConfig {
+    fn default() -> Self {
+        Self {
+            bucket_boundaries: vec![64, 128, 256, 512, 1024, 2048],
+            max_tokens_per_batch: 4096,
+            pad_to_multiple_of: 8,
+            shuffle_within_bucket: true,
+            seed: 42,
+        }
+    }
+}
+
+/// A batch of sequence indices with associated length and efficiency metadata.
+#[derive(Debug, Clone)]
+pub struct LengthBatch {
+    /// Indices into the original dataset (one per sequence in this batch).
+    pub indices: Vec<usize>,
+    /// The effective padded length applied to every sequence in this batch.
+    pub padded_length: usize,
+    /// Number of non-padding tokens across all sequences.
+    pub real_tokens: usize,
+    /// Total tokens including padding: `indices.len() * padded_length`.
+    pub total_tokens: usize,
+}
+
+impl LengthBatch {
+    /// Fraction of tokens that carry real content (1 − padding_rate).
+    ///
+    /// Returns `1.0` for empty batches.
+    #[must_use]
+    pub fn efficiency(&self) -> f32 {
+        if self.total_tokens == 0 {
+            return 1.0;
+        }
+        self.real_tokens as f32 / self.total_tokens as f32
+    }
+}
+
+/// Length-grouped sampler.
+///
+/// Given per-sample sequence lengths, groups samples into fixed-boundary buckets
+/// and packs them into token-budget batches. All sequences inside one batch are
+/// padded to the same length (the bucket ceiling), so padding waste is bounded by
+/// `(bucket_ceiling - min_len_in_bucket) / bucket_ceiling`.
+///
+/// # Design
+///
+/// - **Bucket assignment** is O(n · B) where B = number of boundaries (typically 6).
+/// - **Batch generation** is O(n) per epoch.
+/// - **Shuffling** uses a minimal LCG (Knuth multiplicative) for O(n) Fisher-Yates
+///   with no heap allocation beyond the bucket vectors.
+pub struct LengthGroupedSampler {
+    /// Bucketing configuration.
+    pub config: BucketConfig,
+    /// `lengths[i]` = sequence length of sample `i`.
+    lengths: Vec<usize>,
+    /// `bucket_assignments[i]` = bucket index for sample `i`.
+    bucket_assignments: Vec<usize>,
+    /// Total number of buckets (`config.bucket_boundaries.len() + 1`).
+    num_buckets: usize,
+}
+
+impl LengthGroupedSampler {
+    /// Create a sampler from per-sample lengths.
+    ///
+    /// # Arguments
+    ///
+    /// * `lengths` — one entry per dataset sample; may be empty.
+    /// * `config`  — bucketing parameters.
+    #[must_use]
+    pub fn new(lengths: Vec<usize>, config: BucketConfig) -> Self {
+        let num_buckets = config.bucket_boundaries.len() + 1;
+        let bucket_assignments: Vec<usize> = lengths
+            .iter()
+            .map(|&l| {
+                config
+                    .bucket_boundaries
+                    .iter()
+                    .position(|&b| l <= b)
+                    .unwrap_or(num_buckets - 1)
+            })
+            .collect();
+        Self {
+            config,
+            lengths,
+            bucket_assignments,
+            num_buckets,
+        }
+    }
+
+    /// Return the bucket index for an arbitrary sequence `length`.
+    #[must_use]
+    pub fn bucket_for_length(&self, length: usize) -> usize {
+        self.config
+            .bucket_boundaries
+            .iter()
+            .position(|&b| length <= b)
+            .unwrap_or(self.num_buckets - 1)
+    }
+
+    /// Return the effective upper bound on sequence length for `bucket_id`.
+    ///
+    /// For all buckets except the last this is the configured boundary.
+    /// For the overflow bucket it is the maximum of the sequences actually assigned
+    /// to it, or 0 if the bucket is empty.
+    #[must_use]
+    pub fn bucket_max_length(&self, bucket_id: usize) -> usize {
+        if bucket_id < self.config.bucket_boundaries.len() {
+            self.config.bucket_boundaries[bucket_id]
+        } else {
+            // Overflow bucket: use the actual maximum length.
+            self.lengths
+                .iter()
+                .zip(self.bucket_assignments.iter())
+                .filter(|&(_, &b)| b == bucket_id)
+                .map(|(&l, _)| l)
+                .max()
+                .unwrap_or(0)
+        }
+    }
+
+    /// Round `length` up to the nearest multiple of `pad_to_multiple_of`.
+    ///
+    /// Returns `length` unchanged when `pad_to_multiple_of` is 0 or 1.
+    #[must_use]
+    pub fn padded_length(length: usize, pad_to_multiple_of: usize) -> usize {
+        if pad_to_multiple_of <= 1 {
+            return length;
+        }
+        ((length + pad_to_multiple_of - 1) / pad_to_multiple_of) * pad_to_multiple_of
+    }
+
+    /// Generate all batches for one epoch.
+    ///
+    /// Batches are ordered bucket-by-bucket (shortest → longest). Within each
+    /// bucket the ordering is either the original insertion order or a
+    /// deterministic LCG shuffle depending on [`BucketConfig::shuffle_within_bucket`].
+    ///
+    /// The number of sequences per batch is derived dynamically:
+    /// `batch_size = max_tokens_per_batch / padded_bucket_length`.
+    /// This ensures every batch fits within the token budget while maximising GPU
+    /// occupancy.
+    #[must_use]
+    pub fn generate_batches(&self) -> Vec<LengthBatch> {
+        // Distribute sample indices across buckets.
+        let mut buckets: Vec<Vec<usize>> = vec![Vec::new(); self.num_buckets];
+        for (idx, &bucket_id) in self.bucket_assignments.iter().enumerate() {
+            buckets[bucket_id].push(idx);
+        }
+
+        // Optionally shuffle each bucket in-place with a deterministic LCG
+        // (Knuth multiplicative, 64-bit Marsaglia variant).
+        if self.config.shuffle_within_bucket {
+            let mut state = self.config.seed;
+            for bucket in &mut buckets {
+                // Fisher-Yates using LCG for the index draw.
+                for i in (1..bucket.len()).rev() {
+                    state = state
+                        .wrapping_mul(6_364_136_223_846_793_005)
+                        .wrapping_add(1_442_695_040_888_963_407);
+                    // Use upper 31 bits (higher quality) for the modulo.
+                    let j = ((state >> 33) as usize) % (i + 1);
+                    bucket.swap(i, j);
+                }
+            }
+        }
+
+        let mut all_batches = Vec::new();
+
+        for (bucket_id, bucket_indices) in buckets.iter().enumerate() {
+            if bucket_indices.is_empty() {
+                continue;
+            }
+
+            let bucket_padded = Self::padded_length(
+                self.bucket_max_length(bucket_id),
+                self.config.pad_to_multiple_of,
+            );
+
+            // Guarantee at least one sequence per batch even if a single
+            // sequence already exceeds the token budget.
+            let batch_size = (self.config.max_tokens_per_batch / bucket_padded.max(1)).max(1);
+
+            for chunk in bucket_indices.chunks(batch_size) {
+                let real_tokens: usize = chunk.iter().map(|&i| self.lengths[i]).sum();
+                let total_tokens = chunk.len() * bucket_padded;
+                all_batches.push(LengthBatch {
+                    indices: chunk.to_vec(),
+                    padded_length: bucket_padded,
+                    real_tokens,
+                    total_tokens,
+                });
+            }
+        }
+
+        all_batches
+    }
+
+    /// Overall padding efficiency across all batches for one epoch.
+    ///
+    /// Defined as `real_tokens / total_tokens`. A value of 1.0 means no padding.
+    #[must_use]
+    pub fn padding_efficiency(&self) -> f32 {
+        let batches = self.generate_batches();
+        let real: usize = batches.iter().map(|b| b.real_tokens).sum();
+        let total: usize = batches.iter().map(|b| b.total_tokens).sum();
+        if total == 0 {
+            return 1.0;
+        }
+        real as f32 / total as f32
+    }
+
+    /// Padding efficiency for the naive baseline where every sequence is padded
+    /// to the global maximum length.
+    ///
+    /// Comparing [`Self::padding_efficiency`] against this value quantifies the
+    /// benefit of bucketing on the current dataset.
+    #[must_use]
+    pub fn baseline_efficiency(&self) -> f32 {
+        let max_len = self.lengths.iter().copied().max().unwrap_or(1);
+        let padded_max = Self::padded_length(max_len, self.config.pad_to_multiple_of);
+        let total = self.lengths.len() * padded_max;
+        let real: usize = self.lengths.iter().sum();
+        if total == 0 {
+            return 1.0;
+        }
+        real as f32 / total as f32
+    }
+
+    /// Number of batches produced per epoch.
+    ///
+    /// Equivalent to `generate_batches().len()` but slightly cheaper to express.
+    #[must_use]
+    pub fn num_batches(&self) -> usize {
+        self.generate_batches().len()
+    }
+}
+
+/// Pack a list of sequences into batches greedily by total token count.
+///
+/// Sequences are added to the current batch in order; when adding the next
+/// sequence would push the running total above `max_tokens`, the current batch
+/// is flushed and a new one is started. The returned index vectors reference
+/// positions in the original `lengths` slice.
+///
+/// Unlike [`LengthGroupedSampler`] this function does not sort or reorder
+/// sequences — it respects the caller's ordering, which enables integration
+/// with pre-sorted iterators or interleaved multi-dataset streams.
+///
+/// # Arguments
+///
+/// * `lengths`    — per-sequence lengths in the order they should be packed.
+/// * `max_tokens` — token budget per batch (inclusive upper bound).
+///
+/// # Panics
+///
+/// Does not panic. A sequence whose individual length exceeds `max_tokens` is
+/// placed in a batch by itself (batch total may exceed the budget for that one
+/// batch — callers that cannot tolerate this should pre-filter sequences).
+#[must_use]
+pub fn pack_into_batch(lengths: &[usize], max_tokens: usize) -> Vec<Vec<usize>> {
+    let mut batches: Vec<Vec<usize>> = Vec::new();
+    let mut current_batch: Vec<usize> = Vec::new();
+    let mut current_total: usize = 0;
+
+    for (i, &len) in lengths.iter().enumerate() {
+        // Flush the current batch before adding this sequence — unless the batch
+        // is already empty (handles the oversized-sequence edge case).
+        if !current_batch.is_empty() && current_total + len > max_tokens {
+            batches.push(std::mem::take(&mut current_batch));
+            current_total = 0;
+        }
+        current_batch.push(i);
+        current_total += len;
+    }
+
+    if !current_batch.is_empty() {
+        batches.push(current_batch);
+    }
+
+    batches
+}
+
+/// Compute the padding ratio for naive fixed-size batches.
+///
+/// Within each batch every sequence is padded to the length of the longest
+/// sequence in that batch. The returned value is the fraction of total padded
+/// tokens that are padding (i.e. wasted):
+///
+/// ```text
+/// padding_ratio = 1 − (sum of real tokens) / (sum of padded tokens)
+/// ```
+///
+/// A ratio of 0.0 means perfect packing; 0.5 means half the memory/compute is
+/// wasted on padding.
+///
+/// # Arguments
+///
+/// * `batches` — batch index vectors as produced by [`pack_into_batch`].
+/// * `lengths` — original per-sequence lengths (indexed by the values in `batches`).
+#[must_use]
+pub fn naive_padding_ratio(batches: &[Vec<usize>], lengths: &[usize]) -> f32 {
+    let mut total_real: usize = 0;
+    let mut total_padded: usize = 0;
+
+    for batch in batches {
+        let max_len = batch.iter().map(|&i| lengths[i]).max().unwrap_or(0);
+        for &i in batch {
+            total_real += lengths[i];
+            total_padded += max_len;
+        }
+    }
+
+    if total_padded == 0 {
+        return 0.0;
+    }
+    1.0 - total_real as f32 / total_padded as f32
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_bucket_assignment() {
+        let config = BucketConfig {
+            bucket_boundaries: vec![64, 128],
+            shuffle_within_bucket: false,
+            ..Default::default()
+        };
+        let lengths = vec![32, 64, 65, 128, 129, 500];
+        let sampler = LengthGroupedSampler::new(lengths, config);
+        assert_eq!(sampler.bucket_for_length(32), 0);
+        assert_eq!(sampler.bucket_for_length(64), 0);
+        assert_eq!(sampler.bucket_for_length(65), 1);
+        assert_eq!(sampler.bucket_for_length(128), 1);
+        assert_eq!(sampler.bucket_for_length(129), 2);
+    }
+
+    #[test]
+    fn test_padded_length_multiple() {
+        assert_eq!(LengthGroupedSampler::padded_length(65, 8), 72);
+        assert_eq!(LengthGroupedSampler::padded_length(64, 8), 64);
+        assert_eq!(LengthGroupedSampler::padded_length(0, 8), 0);
+    }
+
+    #[test]
+    fn test_generate_batches_all_indices_covered() {
+        let lengths: Vec<usize> = (1..=100).map(|i| i * 10).collect();
+        let config = BucketConfig {
+            shuffle_within_bucket: false,
+            max_tokens_per_batch: 2048,
+            ..Default::default()
+        };
+        let sampler = LengthGroupedSampler::new(lengths, config);
+        let batches = sampler.generate_batches();
+        let mut seen: Vec<bool> = vec![false; 100];
+        for batch in &batches {
+            for &idx in &batch.indices {
+                seen[idx] = true;
+            }
+        }
+        assert!(seen.iter().all(|&v| v), "all indices must appear in batches");
+    }
+
+    #[test]
+    fn test_bucket_efficiency_beats_baseline() {
+        // Power-law-like lengths: many short, few long.
+        let lengths: Vec<usize> = (0..200)
+            .map(|i| {
+                if i < 150 {
+                    50 + i % 30
+                } else {
+                    500 + i * 5
+                }
+            })
+            .collect();
+        let config = BucketConfig {
+            shuffle_within_bucket: false,
+            max_tokens_per_batch: 4096,
+            ..Default::default()
+        };
+        let sampler = LengthGroupedSampler::new(lengths, config);
+        let bucketed = sampler.padding_efficiency();
+        let baseline = sampler.baseline_efficiency();
+        assert!(
+            bucketed > baseline,
+            "bucketed efficiency {bucketed:.3} must beat baseline {baseline:.3}"
+        );
+    }
+
+    #[test]
+    fn test_batch_total_tokens_correct() {
+        let lengths = vec![64, 64, 64, 64]; // 4 equal sequences
+        let config = BucketConfig {
+            bucket_boundaries: vec![64],
+            shuffle_within_bucket: false,
+            max_tokens_per_batch: 512,
+            pad_to_multiple_of: 1,
+            seed: 0,
+        };
+        let sampler = LengthGroupedSampler::new(lengths, config);
+        let batches = sampler.generate_batches();
+        for batch in &batches {
+            assert_eq!(batch.total_tokens, batch.indices.len() * batch.padded_length);
+        }
+    }
+
+    #[test]
+    fn test_pack_into_batch_respects_max_tokens() {
+        let lengths = vec![100, 200, 150, 300, 50];
+        let batches = pack_into_batch(&lengths, 400);
+        for batch in &batches {
+            // Each individual batch should not exceed max_tokens unless it is a
+            // single oversized sequence (which would be placed alone).
+            let total: usize = batch.iter().map(|&i| lengths[i]).sum();
+            // Only a single-element batch may exceed the limit.
+            if batch.len() > 1 {
+                assert!(
+                    total <= 400,
+                    "multi-element batch total {total} exceeds max 400"
+                );
+            }
+        }
+        // Stricter check: no batch in this test should exceed the limit because
+        // no individual sequence exceeds 400.
+        for batch in &batches {
+            let total: usize = batch.iter().map(|&i| lengths[i]).sum();
+            assert!(total <= 400, "batch total {total} exceeds max 400");
+        }
+    }
+
+    #[test]
+    fn test_pack_into_batch_covers_all() {
+        let lengths = vec![10, 20, 30, 40, 50];
+        let batches = pack_into_batch(&lengths, 100);
+        let mut seen = vec![false; 5];
+        for batch in &batches {
+            for &i in batch {
+                seen[i] = true;
+            }
+        }
+        assert!(seen.iter().all(|&v| v));
+    }
+
+    #[test]
+    fn test_naive_padding_ratio() {
+        // Two batches: [len=10, len=10] and [len=50]
+        let batches = vec![vec![0usize, 1], vec![2]];
+        let lengths = vec![10, 10, 50];
+        let ratio = naive_padding_ratio(&batches, &lengths);
+        // Batch 0: both padded to 10 → 0 waste. Batch 1: padded to 50 → 0 waste.
+        assert!((ratio - 0.0).abs() < 1e-6);
+    }
+
+    #[test]
+    fn test_naive_padding_ratio_with_waste() {
+        // Batch: [len=10, len=100] → padded to 100: real=110, padded=200 → ratio=0.45
+        let batches = vec![vec![0usize, 1]];
+        let lengths = vec![10, 100];
+        let ratio = naive_padding_ratio(&batches, &lengths);
+        assert!((ratio - 0.45).abs() < 0.01);
+    }
+
+    #[test]
+    fn test_num_buckets_overflow() {
+        // Sequences longer than all boundaries go to the last bucket.
+        let lengths = vec![10000, 20000];
+        let config = BucketConfig {
+            bucket_boundaries: vec![64, 128],
+            shuffle_within_bucket: false,
+            max_tokens_per_batch: 100_000,
+            ..Default::default()
+        };
+        let sampler = LengthGroupedSampler::new(lengths, config);
+        let batches = sampler.generate_batches();
+        let total_indices: usize = batches.iter().map(|b| b.indices.len()).sum();
+        assert_eq!(total_indices, 2);
+    }
+
+    #[test]
+    fn test_batch_efficiency() {
+        let batch = LengthBatch {
+            indices: vec![0, 1],
+            padded_length: 100,
+            real_tokens: 150,
+            total_tokens: 200,
+        };
+        assert!((batch.efficiency() - 0.75).abs() < 1e-6);
+    }
+
+    // --- Additional robustness tests ---
+
+    #[test]
+    fn test_empty_dataset() {
+        let sampler =
+            LengthGroupedSampler::new(vec![], BucketConfig::default());
+        assert_eq!(sampler.generate_batches().len(), 0);
+        assert!((sampler.padding_efficiency() - 1.0).abs() < 1e-6);
+        assert!((sampler.baseline_efficiency() - 1.0).abs() < 1e-6);
+    }
+
+    #[test]
+    fn test_single_sequence() {
+        let lengths = vec![37usize];
+        let config = BucketConfig {
+            pad_to_multiple_of: 8,
+            shuffle_within_bucket: false,
+            max_tokens_per_batch: 256,
+            ..Default::default()
+        };
+        let sampler = LengthGroupedSampler::new(lengths, config);
+        let batches = sampler.generate_batches();
+        assert_eq!(batches.len(), 1);
+        // 37 falls in bucket 0 (boundary = 64). padded_length(64, 8) = 64.
+        // Sequences in the same bucket are padded to the bucket ceiling so that
+        // all members of the bucket can share a fixed batch dimension.
+        assert_eq!(batches[0].padded_length, 64);
+        assert_eq!(batches[0].real_tokens, 37);
+        assert_eq!(batches[0].total_tokens, 64);
+    }
+
+    #[test]
+    fn test_uniform_lengths_no_waste() {
+        // When all sequences have the same length, efficiency == 1.0 regardless
+        // of how many buckets there are.
+        let lengths = vec![128usize; 64];
+        let config = BucketConfig {
+            pad_to_multiple_of: 1,
+            shuffle_within_bucket: false,
+            max_tokens_per_batch: 1024,
+            ..Default::default()
+        };
+        let sampler = LengthGroupedSampler::new(lengths, config);
+        assert!((sampler.padding_efficiency() - 1.0).abs() < 1e-6);
+    }
+
+    #[test]
+    fn test_shuffle_produces_all_indices() {
+        // Shuffled batches must still cover every sample exactly once.
+        let lengths: Vec<usize> = (1..=50).map(|i| i * 7).collect();
+        let config = BucketConfig {
+            shuffle_within_bucket: true,
+            max_tokens_per_batch: 1024,
+            seed: 123,
+            ..Default::default()
+        };
+        let sampler = LengthGroupedSampler::new(lengths.clone(), config);
+        let batches = sampler.generate_batches();
+        let mut counts = vec![0usize; lengths.len()];
+        for batch in &batches {
+            for &idx in &batch.indices {
+                counts[idx] += 1;
+            }
+        }
+        assert!(counts.iter().all(|&c| c == 1), "each index must appear exactly once");
+    }
+
+    #[test]
+    fn test_padded_length_no_multiple() {
+        // pad_to_multiple_of = 0 and = 1 both disable alignment.
+        assert_eq!(LengthGroupedSampler::padded_length(65, 0), 65);
+        assert_eq!(LengthGroupedSampler::padded_length(65, 1), 65);
+    }
+
+    #[test]
+    fn test_pack_empty() {
+        let batches = pack_into_batch(&[], 1024);
+        assert!(batches.is_empty());
+    }
+
+    #[test]
+    fn test_pack_oversized_single_sequence() {
+        // A sequence longer than max_tokens must still appear in exactly one batch.
+        let lengths = vec![2048usize];
+        let batches = pack_into_batch(&lengths, 512);
+        assert_eq!(batches.len(), 1);
+        assert_eq!(batches[0], vec![0]);
+    }
+}

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

@@ -2,6 +2,7 @@
 
 pub mod comprehensive_integration_test;
 pub mod end_to_end_training_example;
+pub mod length_bucketing;
 pub mod gradient_accumulation_enhanced;
 pub mod gradient_accumulator;
 pub mod gradient_clipping;
@@ -41,6 +42,9 @@ pub use transformer_trainer::{
     ClassMetrics, EvaluationResults, ModelConfig, ModelOutput, TrainingMetrics, TransformerModel,
     TransformerTrainer,
 };
+pub use length_bucketing::{
+    BucketConfig, LengthBatch, LengthGroupedSampler, naive_padding_ratio, pack_into_batch,
+};
 
 /// Training state structure
 #[derive(Debug, Clone)]