ferrox_core/
alibi.rs

1//! **ALiBi** -- the per-head linear position bias llama.cpp adds inside
2//! `ggml_soft_max_ext` when `hparams.f_max_alibi_bias > 0`.
3//!
4//! # What it is
5//!
6//! `ggml_soft_max_ext(kq, kq_mask, scale, max_bias)` computes
7//! `softmax(scale * kq + slope_h * mask)`, where the mask a model with
8//! `use_alibi` fills is `-|p_key - p_query|` on the visible entries and
9//! `-inf` elsewhere (`llama-kv-cache.cpp:1673-1676`), and the slope is
10//! `ggml-cpu/ops.cpp:5489-5508`:
11//!
12//! ```text
13//! n_head_log2 = 2^floor(log2(n_head))
14//! m0 = 2^(-max_bias / n_head_log2)
15//! m1 = 2^(-(max_bias / 2) / n_head_log2)
16//! slope_h = h < n_head_log2 ? m0^(h + 1) : m1^(2 (h - n_head_log2) + 1)
17//! ```
18//!
19//! So a query at `p_q` sees key `p_k <= p_q` with `slope_h * (p_k -
20//! p_q)` added to its scaled score, after any softcap (the softcapped
21//! graphs scale and `tanh` BEFORE `ggml_soft_max_ext`; none of them
22//! uses ALiBi, but the order is the graph's). [`slopes`] is the
23//! formula; the kernels in `attention` take the result as an optional
24//! per-head slice and rotate nothing for such a model.
25//!
26//! The graphs that set `f_max_alibi_bias` are the caller's business
27//! (`ferrox_models::alibi`); this module is the arithmetic.
28
29/// One slope per query head, `n_heads` long, for a `max_bias` that is
30/// positive. Returns `None` for a non-positive `max_bias`, which is
31/// llama.cpp's "no ALiBi" (`slope = 1.0` on a mask that is then
32/// `0 / -inf`, i.e. no bias at all).
33pub fn slopes(n_heads: usize, max_bias: f32) -> Option<Vec<f32>> {
34    // `max_bias <= 0.0` is false for NaN, and NaN is "no ALiBi" too.
35    if max_bias.partial_cmp(&0.0) != Some(std::cmp::Ordering::Greater) || n_heads == 0 {
36        return None;
37    }
38    let n_head_log2 = 1usize << (usize::BITS - 1 - n_heads.leading_zeros());
39    let m0 = 2f32.powf(-max_bias / n_head_log2 as f32);
40    let m1 = 2f32.powf(-(max_bias / 2.0) / n_head_log2 as f32);
41    Some(
42        (0..n_heads)
43            .map(|h| {
44                if h < n_head_log2 {
45                    m0.powi(h as i32 + 1)
46                } else {
47                    m1.powi(2 * (h - n_head_log2) as i32 + 1)
48                }
49            })
50            .collect(),
51    )
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    /// The textbook case: a power-of-two head count and `max_bias 8`
59    /// gives `2^(-8 (h+1) / n_head)`, i.e. `1/2, 1/4, ..., 1/256` for
60    /// eight heads.
61    #[test]
62    fn a_power_of_two_head_count_is_the_geometric_sequence() {
63        let s = slopes(8, 8.0).unwrap();
64        let want = [
65            0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625, 0.0078125, 0.00390625,
66        ];
67        for (g, w) in s.iter().zip(want) {
68            assert!((g - w).abs() < 1e-7, "{g} vs {w}");
69        }
70    }
71
72    /// The non-power-of-two rule llama.cpp copies from the paper: the
73    /// first `n_head_log2` heads take `m0`, the rest interleave `m1` at
74    /// odd powers. Twelve heads (GPT-2-small's count, MPT-7B is 32).
75    #[test]
76    fn a_non_power_of_two_head_count_takes_the_second_base_for_the_tail() {
77        let s = slopes(12, 8.0).unwrap();
78        // n_head_log2 = 8: m0 = 2^-1, m1 = 2^-0.5.
79        let m1 = 2f32.powf(-0.5);
80        for (h, got) in s.iter().enumerate() {
81            let want = if h < 8 {
82                0.5f32.powi(h as i32 + 1)
83            } else {
84                m1.powi(2 * (h as i32 - 8) + 1)
85            };
86            assert!((got - want).abs() < 1e-7, "head {h}: {got} vs {want}");
87        }
88    }
89
90    #[test]
91    fn a_zero_max_bias_is_no_alibi() {
92        assert!(slopes(8, 0.0).is_none());
93        assert!(slopes(8, -1.0).is_none());
94        assert!(slopes(0, 8.0).is_none());
95    }
96}
ferrox_core/alibi.rs

ferrox_core/
alibi.rs
