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/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");
+        }
+    }
+}