burn_mamba/mamba3/
helpers.rs

1//! Shared helpers used by both [`Mamba3::forward`](super::mamba3::Mamba3::forward)
2//! and [`Mamba3::step`](super::mamba3::Mamba3::step). They isolate three blocks
3//! that previously appeared in both methods at different ranks:
4//!
5//! 1. Trapezoidal discretisation: `dt`, `α`, `β`, `γ`, `da`.
6//! 2. QK-norm + GQA expansion + per-(head, mimo-rank) bias on B / C.
7//! 3. MIMO `V` construction: broadcast-multiply `x` by `mimo_x_hmp`.
8//!
9//! Each helper is generic over the rank `D` of the data tensors so a single
10//! definition serves both the sequence-aware (`forward`) and single-token
11//! (`step`) code paths.
12
13use crate::modules::RmsNorm;
14use crate::modules::gqa_expand_to_heads;
15use crate::modules::softplus;
16use burn::prelude::*;
17
18/// Output of [`trapezoidal_coefficients`].
19///
20/// All tensors share the rank `D` of the inputs.
21pub struct TrapezoidCoeffs<const D: usize> {
22    /// `Δₜ = softplus(dd_dt + dt_bias)`, clamped.
23    pub dt: Tensor<D>,
24    /// `Δₜ · Aₜ` (negative; the log-decay).
25    pub da: Tensor<D>,
26    /// `αₜ = exp(Δₜ · Aₜ) ∈ (0, 1]` — decay.
27    pub alpha: Tensor<D>,
28    /// `βₜ = (1 − λₜ) · Δₜ · αₜ` — left-endpoint weight.
29    pub beta: Tensor<D>,
30    /// `γₜ = λₜ · Δₜ` — right-endpoint weight.
31    pub gamma: Tensor<D>,
32}
33
34/// Compute the trapezoidal discretisation coefficients from the raw
35/// (data-dependent) projections. See the top-of-`mamba3.rs` docs for the
36/// formulas.
37///
38/// All four data tensors share rank `D` and have `nheads` as the last dim.
39/// `dt_bias_h` is broadcast to match.
40pub fn trapezoidal_coefficients<const D: usize>(
41    dd_dt: Tensor<D>,
42    dd_a_raw: Tensor<D>,
43    lambda_raw: Tensor<D>,
44    dt_bias_h: Tensor<1>,
45    dt_limit: (f64, f64),
46    a_floor: f64,
47) -> TrapezoidCoeffs<D> {
48    // Broadcast dt_bias_h [nheads] → [1, ..., 1, nheads] so the addition aligns
49    // on the last dim regardless of leading shape.
50    let dt_bias_broadcast = dt_bias_h.unsqueeze::<D>();
51    let dt = softplus(dd_dt + dt_bias_broadcast).clamp(dt_limit.0, dt_limit.1);
52    // `A = −max(softplus(·), a_floor) ∈ (−∞, −a_floor]`. The floor must be
53    // applied to the (positive) softplus *before* negating: a method call
54    // binds tighter than unary minus, so `-softplus(x).clamp(NEG_INFINITY,
55    // -a_floor)` would collapse the positive softplus to the constant
56    // `-a_floor` and yield `A ≡ +a_floor` — a *growing* state (`α > 1`) with a
57    // dead `dd_A` projection.
58    let a = -softplus(dd_a_raw).clamp(a_floor, f64::INFINITY);
59    let da = dt.clone() * a;
60    let lambda = burn::tensor::activation::sigmoid(lambda_raw);
61    let alpha = da.clone().exp();
62    let beta = (-lambda.clone() + 1.0) * dt.clone() * alpha.clone();
63    let gamma = lambda * dt.clone();
64    TrapezoidCoeffs {
65        dt,
66        da,
67        alpha,
68        beta,
69        gamma,
70    }
71}
72
73/// QK-Norm → GQA-expand groups→heads → add per-(head, mimo-rank) bias.
74///
75/// The input is the raw B/C projection already reshaped to expose the group
76/// dim, with last dim = `state_rank`. The output replaces the group dim with
77/// the head dim, leaving the last dim untouched.
78///
79/// `DP1 = D + 1` (required by [`gqa_expand_to_heads`]'s intermediate rank).
80pub fn qk_norm_expand_bias<const D: usize, const DP1: usize>(
81    raw_mgr: Tensor<D>,
82    norm: &RmsNorm,
83    bias_hmr: Tensor<3>,
84    group_dim: usize,
85    nheads: usize,
86) -> Tensor<D> {
87    // RmsNorm operates on the last dim only, so the leading shape passes through.
88    let normed = norm.forward(raw_mgr);
89    let expanded = gqa_expand_to_heads::<D, DP1>(normed, group_dim, nheads);
90    // Broadcast bias [nheads, mimo_rank, state_rank] → [1, ..., 1, mimo_rank, nheads, state_rank].
91    let bias = bias_hmr.permute([1, 0, 2]).unsqueeze::<D>();
92    expanded + bias
93}
94
95/// Build the MIMO value tensor `v = x ⊙ mimo_x` with broadcasting.
96///
97/// Inserts a `mimo_rank` axis at `insert_dim`. When `mimo_x_hmp` is `None`
98/// (SISO), the inserted axis has size 1 and `x` is passed through; otherwise
99/// broadcasting fills the inserted axis to size `mimo_rank`.
100///
101/// `DP1 = D + 1`.
102pub fn build_v_with_mimo<const D: usize, const DP1: usize>(
103    x: Tensor<D>,
104    mimo_x_hmp: Option<&Tensor<3>>,
105    insert_dim: usize,
106) -> Tensor<DP1> {
107    let x_with_rank_axis = x.unsqueeze_dim::<DP1>(insert_dim);
108    match mimo_x_hmp {
109        None => x_with_rank_axis,
110        Some(mimo_x_hmp) => {
111            // mimo_x_hmp [nheads, mimo_rank, per_head_dim] → permute to
112            // [mimo_rank, nheads, per_head_dim] → unsqueeze leading 1s. The
113            // result broadcasts against `x_with_rank_axis` over (batch, seq, …).
114            let mimo_x_broadcast = mimo_x_hmp.clone().permute([1, 0, 2]).unsqueeze::<DP1>();
115            x_with_rank_axis * mimo_x_broadcast
116        }
117    }
118}
burn_mamba/mamba3/helpers.rs

burn_mamba/mamba3/
helpers.rs
