burn_dragon_kernel 0.5.0

Fused GPU kernel crate for burn_dragon execution paths
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

#[derive(Clone, Copy, Debug)]
pub struct MambaBenchCase {
    pub batch: usize,
    pub time: usize,
    pub d_model: usize,
    pub d_state: usize,
    pub d_conv: usize,
    pub expand: usize,
}

pub const LARGE_RUNG_CASE: MambaBenchCase = MambaBenchCase {
    batch: 1,
    time: 256,
    d_model: 1024,
    d_state: 16,
    d_conv: 4,
    expand: 2,
};