avarok_core/kimi_k3/
ops.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Tiny host GEMV / gather helpers for the C1 CPU graph.
4
5/// `y = W x` with `W` row-major `[out, inn]`.
6pub fn matvec(w: &[f32], x: &[f32], out: usize, inn: usize) -> Vec<f32> {
7    assert_eq!(
8        w.len(),
9        out * inn,
10        "matvec weight {} vs {out}x{inn}",
11        w.len()
12    );
13    assert_eq!(x.len(), inn);
14    let mut y = vec![0.0f32; out];
15    // Interleave independent output rows so the CPU can overlap accumulations.
16    // Each row retains the scalar reduction order: no reassociation or FMA.
17    let grouped = out / 4 * 4;
18    for o in (0..grouped).step_by(4) {
19        let a = &w[o * inn..(o + 1) * inn];
20        let b = &w[(o + 1) * inn..(o + 2) * inn];
21        let c = &w[(o + 2) * inn..(o + 3) * inn];
22        let d = &w[(o + 3) * inn..(o + 4) * inn];
23        let mut acc = [0.0f32; 4];
24        for i in 0..inn {
25            let value = x[i];
26            acc[0] += a[i] * value;
27            acc[1] += b[i] * value;
28            acc[2] += c[i] * value;
29            acc[3] += d[i] * value;
30        }
31        y[o..o + 4].copy_from_slice(&acc);
32    }
33    for o in grouped..out {
34        let row = &w[o * inn..(o + 1) * inn];
35        let mut acc = 0.0f32;
36        for i in 0..inn {
37            acc += row[i] * x[i];
38        }
39        y[o] = acc;
40    }
41    y
42}
43
44/// Column-parallel `y = W x` (Megatron `o_proj`). `W` is `[out, inn]`.
45///
46/// `world == 1` is the unsplit GEMV. `world == 2` splits columns across two
47/// in-process ranks and allreduces (sum) the hidden. `drop_rank` zeros that
48/// shard (C7 known-bad: drop rank 1).
49pub fn matvec_column_tp(
50    w: &[f32],
51    x: &[f32],
52    out: usize,
53    inn: usize,
54    world: usize,
55    drop_rank: Option<usize>,
56) -> Vec<f32> {
57    match world {
58        1 => {
59            assert!(
60                drop_rank.is_none(),
61                "C7 drop_rank requires TP=2, got {drop_rank:?}"
62            );
63            matvec(w, x, out, inn)
64        }
65        2 => matvec_column_tp2(w, x, out, inn, drop_rank),
66        other => panic!("C7 dummy implements TP=1/2, got {other}"),
67    }
68}
69
70/// Two-rank in-process column split + sum. Sequential ranks; the add is the
71/// allreduce. Identity `W` is f32 bit-exact vs [`matvec`].
72fn matvec_column_tp2(
73    w: &[f32],
74    x: &[f32],
75    out: usize,
76    inn: usize,
77    drop_rank: Option<usize>,
78) -> Vec<f32> {
79    assert_eq!(w.len(), out * inn, "TP=2 weight {} vs {out}x{inn}", w.len());
80    assert_eq!(x.len(), inn);
81    assert!(
82        inn.is_multiple_of(2),
83        "TP=2 o_proj inner dim must be even, got {inn}"
84    );
85    let half = inn / 2;
86    let shard = |rank: usize| -> Vec<f32> {
87        if drop_rank == Some(rank) {
88            return vec![0.0f32; out];
89        }
90        let start = rank * half;
91        let mut y = vec![0.0f32; out];
92        for o in 0..out {
93            let row = &w[o * inn + start..o * inn + start + half];
94            let xr = &x[start..start + half];
95            let mut acc = 0.0f32;
96            for i in 0..half {
97                acc += row[i] * xr[i];
98            }
99            y[o] = acc;
100        }
101        y
102    };
103    // In-process ranks (thread), then sum. Sequential join is the reduce.
104    std::thread::scope(|s| {
105        let t0 = s.spawn(|| shard(0));
106        let t1 = s.spawn(|| shard(1));
107        let y0 = t0.join().expect("C7 rank 0");
108        let y1 = t1.join().expect("C7 rank 1");
109        y0.iter().zip(y1).map(|(a, b)| a + b).collect()
110    })
111}
112
113/// Token embedding gather: `embed[token, :]`.
114pub fn embed_token(table: &[f32], token: u32, hidden: usize, vocab: usize) -> Vec<f32> {
115    let t = token as usize;
116    assert!(t < vocab, "token {token} >= vocab {vocab}");
117    table[t * hidden..(t + 1) * hidden].to_vec()
118}
119
120/// Greedy next-token: argmax, lower id on ties.
121pub fn argmax(logits: &[f32]) -> u32 {
122    assert!(!logits.is_empty());
123    let mut best_i = 0usize;
124    let mut best = logits[0];
125    for (i, &v) in logits.iter().enumerate().skip(1) {
126        if v > best {
127            best = v;
128            best_i = i;
129        }
130    }
131    best_i as u32
132}
133
134/// Deterministic fill in `(-scale/2, scale/2)`.
135pub fn fill(n: usize, seed: u32, scale: f32) -> Vec<f32> {
136    (0..n)
137        .map(|i| {
138            let h = seed
139                .wrapping_mul(1664525)
140                .wrapping_add(1013904223)
141                .wrapping_add((i as u32).wrapping_mul(2654435761));
142            ((h % 1000) as f32 / 1000.0 - 0.5) * scale
143        })
144        .collect()
145}
146
147/// Ones vector (RMSNorm / identity-ish scales).
148pub fn ones(n: usize) -> Vec<f32> {
149    vec![1.0; n]
150}
151
152/// Row-major identity padded to `[out, inn]`.
153pub fn ident(out: usize, inn: usize) -> Vec<f32> {
154    let mut w = vec![0.0f32; out * inn];
155    for i in 0..out.min(inn) {
156        w[i * inn + i] = 1.0;
157    }
158    w
159}
160
161#[cfg(test)]
162#[path = "ops_matvec_tests.rs"]
163mod matvec_tests;
avarok_core/kimi_k3/ops.rs

avarok_core/kimi_k3/
ops.rs
