brainharmony 0.1.0

Brain-Harmony multimodal brain foundation model — inference in Rust with Burn ML
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

/// Trace per-operation overhead in the encoder pipeline.
/// Compares: sum-of-parts vs actual full forward.
use std::time::Instant;
use burn::prelude::*;

#[cfg(all(feature = "wgpu-f16", not(feature = "wgpu")))]
mod backend {
    pub type B = burn::backend::wgpu::Wgpu<half::f16, i32, u32>;
    pub fn device() -> burn::backend::wgpu::WgpuDevice { burn::backend::wgpu::WgpuDevice::DefaultDevice }
}
#[cfg(all(feature = "wgpu", not(feature = "wgpu-f16")))]
mod backend {
    pub use burn::backend::{Wgpu as B, wgpu::WgpuDevice};
    pub fn device() -> WgpuDevice { WgpuDevice::DefaultDevice }
}
#[cfg(not(any(feature = "wgpu", feature = "wgpu-f16")))]
mod backend {
    pub use burn::backend::NdArray as B;
    pub fn device() -> burn::backend::ndarray::NdArrayDevice { burn::backend::ndarray::NdArrayDevice::Cpu }
}
use backend::{B, device};

fn main() {
    let d = device();
    brainharmony::init_threads(None);

    let seq = 7200usize;
    let embed = 768usize;
    let heads = 12usize;
    let dh = embed / heads;
    let warmup = 3;
    let runs = 5;

    println!("=== Overhead analysis ===\n");

    // 1. Measure QKV split overhead (reshape + narrow + swap_dims)
    let qkv: Tensor<B, 3> = Tensor::random([1, seq, 3 * embed],
        burn::tensor::Distribution::Normal(0.0, 1.0), &d);

    // Warmup
    for _ in 0..warmup {
        let qkv2 = qkv.clone().reshape([1, seq, 3, heads, dh]);
        let q = qkv2.clone().narrow(2, 0, 1).reshape([1, seq, heads, dh]).swap_dims(1, 2);
        let k = qkv2.clone().narrow(2, 1, 1).reshape([1, seq, heads, dh]).swap_dims(1, 2);
        let v = qkv2.narrow(2, 2, 1).reshape([1, seq, heads, dh]).swap_dims(1, 2);
        let _ = (q + k + v).into_data();
    }

    // Method A: current (narrow along dim 2, then reshape+swap)
    let mut times_a = Vec::new();
    for _ in 0..runs {
        let t0 = Instant::now();
        let qkv2 = qkv.clone().reshape([1, seq, 3, heads, dh]);
        let q = qkv2.clone().narrow(2, 0, 1).reshape([1, seq, heads, dh]).swap_dims(1, 2);
        let k = qkv2.clone().narrow(2, 1, 1).reshape([1, seq, heads, dh]).swap_dims(1, 2);
        let v = qkv2.narrow(2, 2, 1).reshape([1, seq, heads, dh]).swap_dims(1, 2);
        let _ = (q + k + v).into_data();
        times_a.push(t0.elapsed().as_secs_f64() * 1000.0);
    }
    let best_a = times_a.iter().cloned().fold(f64::INFINITY, f64::min);

    // Method B: reshape to [B, N, H, 3*Dh], swap, then narrow along last dim
    for _ in 0..warmup {
        let qkv2 = qkv.clone().reshape([1, seq, heads, 3 * dh]).swap_dims(1, 2);
        let q = qkv2.clone().narrow(3, 0, dh);
        let k = qkv2.clone().narrow(3, dh, dh);
        let v = qkv2.narrow(3, 2 * dh, dh);
        let _ = (q + k + v).into_data();
    }
    let mut times_b = Vec::new();
    for _ in 0..runs {
        let t0 = Instant::now();
        let qkv2 = qkv.clone().reshape([1, seq, heads, 3 * dh]).swap_dims(1, 2);
        let q = qkv2.clone().narrow(3, 0, dh);
        let k = qkv2.clone().narrow(3, dh, dh);
        let v = qkv2.narrow(3, 2 * dh, dh);
        let _ = (q + k + v).into_data();
        times_b.push(t0.elapsed().as_secs_f64() * 1000.0);
    }
    let best_b = times_b.iter().cloned().fold(f64::INFINITY, f64::min);

    // Method C: three separate projections (no narrow/split at all)
    let wq: Tensor<B, 2> = Tensor::random([embed, embed], burn::tensor::Distribution::Normal(0.0, 0.01), &d);
    let wk = wq.clone();
    let wv = wq.clone();
    let x2: Tensor<B, 2> = Tensor::random([seq, embed], burn::tensor::Distribution::Normal(0.0, 1.0), &d);

    for _ in 0..warmup {
        let q = x2.clone().matmul(wq.clone()).reshape([1, heads, seq, dh]);
        let k = x2.clone().matmul(wk.clone()).reshape([1, heads, seq, dh]);
        let v = x2.clone().matmul(wv.clone()).reshape([1, heads, seq, dh]);
        let _ = (q + k + v).into_data();
    }
    let mut times_c = Vec::new();
    for _ in 0..runs {
        let t0 = Instant::now();
        let q = x2.clone().matmul(wq.clone()).reshape([1, heads, seq, dh]);
        let k = x2.clone().matmul(wk.clone()).reshape([1, heads, seq, dh]);
        let v = x2.clone().matmul(wv.clone()).reshape([1, heads, seq, dh]);
        let _ = (q + k + v).into_data();
        times_c.push(t0.elapsed().as_secs_f64() * 1000.0);
    }
    let best_c = times_c.iter().cloned().fold(f64::INFINITY, f64::min);

    println!("QKV split methods:");
    println!("  A: narrow(dim2) + reshape + swap  : {best_a:.1}ms");
    println!("  B: reshape + swap + narrow(dim3)  : {best_b:.1}ms");
    println!("  C: 3 separate matmuls + reshape   : {best_c:.1}ms");

    // 2. Measure transpose overhead
    let k4: Tensor<B, 4> = Tensor::random([1, heads, seq, dh],
        burn::tensor::Distribution::Normal(0.0, 1.0), &d);
    for _ in 0..warmup {
        let _ = k4.clone().transpose().into_data();
    }
    let mut times_t = Vec::new();
    for _ in 0..runs {
        let t0 = Instant::now();
        let _ = k4.clone().transpose().into_data();
        times_t.push(t0.elapsed().as_secs_f64() * 1000.0);
    }
    let best_t = times_t.iter().cloned().fold(f64::INFINITY, f64::min);
    println!("\nTranspose [1,12,7200,64] -> [1,12,64,7200]: {best_t:.1}ms");

    // 3. Measure full 1 block vs sum of parts
    let block = brainharmony::model::block::Block::<B>::new(embed, heads, 4.0, true, 1e-6, &d);
    let x3: Tensor<B, 3> = Tensor::random([1, seq, embed],
        burn::tensor::Distribution::Normal(0.0, 1.0), &d);

    for _ in 0..warmup {
        let _ = block.forward(x3.clone(), None).into_data();
    }
    let mut times_block = Vec::new();
    for _ in 0..runs {
        let t0 = Instant::now();
        let out = block.forward(x3.clone(), None);
        let _ = out.into_data();
        times_block.push(t0.elapsed().as_secs_f64() * 1000.0);
    }
    let best_block = times_block.iter().cloned().fold(f64::INFINITY, f64::min);
    println!("\nFull Block forward: {best_block:.1}ms (x12 = {:.0}ms)", best_block * 12.0);
    println!("  Per-op estimate:  ~300ms  (x12 = ~3600ms)");
    println!("  Overhead per block: {:.1}ms", best_block - 300.0);
}