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
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
145

use std::marker::PhantomData;

use cubecl::ir::DeviceProperties;
use cubek_matmul::{
    components::{
        global::{
            GlobalReaderConfig, GlobalWriterConfig, InputLoadFlow, PartitionedStageFamily,
            PlaneFlowConfig, PlaneFlowPartitionRule,
            memory::{GlobalMemoryConfig, ViewDirection},
            multi_stage::EventLoadingMode,
            read::ReaderMode,
        },
        stage::StridedStageFamily,
    },
    definition::LoadingPrecomputeStrategy,
};
use cubek_std::{MatrixLayout, StageIdent};

use crate::{
    components::{
        global::{
            GlobalAttentionFamily,
            simple::{SimpleGlobalAttention, config::SimpleGlobalAttentionConfig},
        },
        stage::{StageAttentionConfig as _, StageAttentionFamily},
    },
    definition::{AttentionBlueprint, AttentionElems, AttentionPrecision, AttentionSetupError},
};

pub struct SimpleGlobalAttentionFamily<SA: StageAttentionFamily> {
    _phantom: PhantomData<SA>,
}

impl<
    SA: StageAttentionFamily<
            KeyStage = StridedStageFamily,
            ValueStage = StridedStageFamily,
            OutStage = PartitionedStageFamily,
        >,
> GlobalAttentionFamily for SimpleGlobalAttentionFamily<SA>
{
    type Attention<AP: AttentionPrecision> = SimpleGlobalAttention<AP, SA::Attention<AP>>;

    type Config = SimpleGlobalAttentionConfig<SA::Config>;

    fn expand_config(
        device_props: &DeviceProperties,
        blueprint: &AttentionBlueprint,
        dtypes: &AttentionElems,
    ) -> Result<Self::Config, AttentionSetupError> {
        let stage_config = SA::expand_config(device_props, blueprint, dtypes)?;

        let precompute_job = LoadingPrecomputeStrategy::Never.into();
        let plane_dim = stage_config.plane_dim();
        let reader_mode = ReaderMode::Relaxed;
        let event_loading_mode = EventLoadingMode::Relaxed;
        let specialization_tensor_config = InputLoadFlow::MainOnly;
        let plane_flow_config = PlaneFlowConfig::new_unspecialized(stage_config.num_planes());

        let query_gmem_config = GlobalMemoryConfig {
            vector_size: blueprint.vector_sizes.query,
            check_row_bounds: blueprint.check_bounds.seq_q,
            check_col_bounds: blueprint.check_bounds.head_dim,
            matrix_layout: MatrixLayout::RowMajor,
            view_direction: ViewDirection::None,
            dtype: dtypes.query_global,
        };

        let mask_gmem_config = GlobalMemoryConfig {
            vector_size: blueprint.vector_sizes.mask,
            check_row_bounds: blueprint.check_bounds.seq_q,
            check_col_bounds: blueprint.check_bounds.seq_kv,
            matrix_layout: MatrixLayout::RowMajor,
            view_direction: ViewDirection::Col,
            dtype: dtypes.mask,
        };

        let key_gmem_config = GlobalMemoryConfig {
            vector_size: blueprint.vector_sizes.key,
            check_row_bounds: blueprint.check_bounds.seq_kv,
            check_col_bounds: blueprint.check_bounds.head_dim,
            matrix_layout: MatrixLayout::RowMajor,
            view_direction: ViewDirection::Row,
            dtype: dtypes.key_global,
        };

        let value_gmem_config = GlobalMemoryConfig {
            vector_size: blueprint.vector_sizes.value,
            check_row_bounds: blueprint.check_bounds.seq_kv,
            check_col_bounds: blueprint.check_bounds.val_dim,
            matrix_layout: MatrixLayout::RowMajor,
            view_direction: ViewDirection::Row,
            dtype: dtypes.value_global,
        };

        let out_gmem_config = GlobalMemoryConfig {
            vector_size: blueprint.vector_sizes.out,
            check_row_bounds: blueprint.check_bounds.seq_q,
            check_col_bounds: blueprint.check_bounds.val_dim,
            matrix_layout: MatrixLayout::RowMajor,
            view_direction: ViewDirection::None,
            dtype: dtypes.out_global,
        };

        let key_reader_config = GlobalReaderConfig {
            gmem_config: key_gmem_config,
            smem_config: stage_config.key_smem_config(),
            precompute_job,
            plane_dim,
            reader_mode,
            event_loading_mode,
            input_load_flow: specialization_tensor_config,
            plane_flow_config,
            stage_ident: StageIdent::Rhs,
        };

        let value_reader_config = GlobalReaderConfig {
            gmem_config: value_gmem_config,
            smem_config: stage_config.value_smem_config(),
            precompute_job,
            plane_dim,
            reader_mode,
            event_loading_mode,
            input_load_flow: specialization_tensor_config,
            plane_flow_config,
            stage_ident: StageIdent::Rhs,
        };

        let writer_config = GlobalWriterConfig {
            gmem_config: out_gmem_config,
            smem_config: stage_config.out_smem_config(),
            plane_flow_partition_rule: PlaneFlowPartitionRule::MainFlowOnly,
            plane_dim,
        };

        Ok(SimpleGlobalAttentionConfig {
            stage_config,
            key_reader_config,
            value_reader_config,
            query_gmem_config,
            mask_gmem_config,
            writer_config,
        })
    }
}