cubek-attention 0.2.0

CubeK: Attention Kernels
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

use cubecl;
use cubecl::prelude::*;

use cubek_std::tile::{Plane, Tile};

#[derive(CubeType)]
/// Query input to the Tile Attention
pub struct Query<L: Numeric> {
    pub tile: Tile<L, Plane, ReadWrite>,
}

#[cube]
impl<L: Numeric> Query<L> {
    pub fn new(tile: Tile<L, Plane, ReadWrite>) -> Query<L> {
        Query::<L> { tile }
    }
}

#[derive(CubeType)]
pub struct Key<R: Numeric> {
    pub tile: Tile<R, Plane, ReadWrite>,
}

#[cube]
impl<R: Numeric> Key<R> {
    pub fn new(tile: Tile<R, Plane, ReadWrite>) -> Key<R> {
        Key::<R> { tile }
    }
}

#[derive(CubeType)]
pub struct Value<R: Numeric> {
    pub tile: Tile<R, Plane, ReadWrite>,
}

#[cube]
impl<R: Numeric> Value<R> {
    pub fn new(tile: Tile<R, Plane, ReadWrite>) -> Value<R> {
        Value::<R> { tile }
    }
}