rudnn/attention/kernel_ir/definition/
spec.rs

1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use half::{bf16, f16};
4
5use crate::attention::kernel_ir::{
6    definition::{AccumulatorPrecision, AttentionGlobalTypes},
7    launch::{AttentionArgs, TensorArgs},
8};
9
10/// Attention spec defining each element types used in the computation as well as
11/// how the arguments are passed to the kernel.
12pub trait AttentionSpec: Send + Sync + Clone + 'static {
13    type Precision: AttentionPrecision;
14    /// How the input and output tensors are passed as arguments.
15    type Args: AttentionArgs;
16}
17
18impl<AP: AttentionPrecision, Args: AttentionArgs> AttentionSpec for (AP, Args) {
19    type Precision = AP;
20    type Args = Args;
21}
22
23// A simple default for TensorArgs
24impl<AP: AttentionPrecision> AttentionSpec for AP {
25    type Precision = AP;
26    type Args = TensorArgs;
27}
28
29pub trait QueryPrecision: Send + Sync + Copy + 'static {
30    type Global: Float;
31    type GlobalSize: Size;
32    type Tile: Float;
33}
34
35pub trait StagedMatrixPrecision: Send + Sync + Copy + 'static {
36    type Global: Float;
37    type GlobalSize: Size;
38    type Stage: Float;
39    type StageSize: Size;
40}
41
42pub trait AttentionPrecision: Send + Sync + Copy + 'static {
43    type Query: QueryPrecision;
44    type Key: StagedMatrixPrecision;
45    type Value: StagedMatrixPrecision;
46    type KVTile: Float;
47    type SoftmaxAcc: Float;
48    type SoftmaxLhs: Float;
49    type Accumulator: Float;
50    type Mask: Numeric;
51    type MaskSize: Size;
52    type Out: StagedMatrixPrecision;
53}
54
55impl QueryPrecision for f16 {
56    type Global = f16;
57    type GlobalSize = Const<0>;
58    type Tile = f16;
59}
60
61impl QueryPrecision for bf16 {
62    type Global = bf16;
63    type GlobalSize = Const<0>;
64    type Tile = bf16;
65}
66
67impl QueryPrecision for flex32 {
68    type Global = f32;
69    type GlobalSize = Const<0>;
70    type Tile = f16;
71}
72
73impl QueryPrecision for f32 {
74    type Global = f32;
75    type GlobalSize = Const<0>;
76    type Tile = f32;
77}
78
79impl QueryPrecision for f64 {
80    type Global = f64;
81    type GlobalSize = Const<0>;
82    type Tile = f32;
83}
84
85impl<G: Float, GS: Size, T: Float> QueryPrecision for (G, GS, T) {
86    type Global = G;
87    type GlobalSize = GS;
88    type Tile = T;
89}
90
91impl StagedMatrixPrecision for f16 {
92    type Global = f16;
93    type GlobalSize = Const<0>;
94    type Stage = f16;
95    type StageSize = Const<0>;
96}
97
98impl StagedMatrixPrecision for bf16 {
99    type Global = bf16;
100    type GlobalSize = Const<0>;
101    type Stage = bf16;
102    type StageSize = Const<0>;
103}
104
105impl StagedMatrixPrecision for flex32 {
106    type Global = f32;
107    type GlobalSize = Const<0>;
108    type Stage = f16;
109    type StageSize = Const<0>;
110}
111
112impl StagedMatrixPrecision for f32 {
113    type Global = f32;
114    type GlobalSize = Const<0>;
115    type Stage = f32;
116    type StageSize = Const<0>;
117}
118
119impl StagedMatrixPrecision for f64 {
120    type Global = f64;
121    type GlobalSize = Const<0>;
122    type Stage = f32;
123    type StageSize = Const<0>;
124}
125
126impl<G: Float, GS: Size, S: Float, SS: Size> StagedMatrixPrecision for (G, GS, S, SS) {
127    type Global = G;
128    type GlobalSize = GS;
129    type Stage = S;
130    type StageSize = SS;
131}
132
133impl AttentionPrecision for f16 {
134    type Query = f16;
135    type Key = f16;
136    type Value = f16;
137    type KVTile = f16;
138    type SoftmaxLhs = f16;
139    #[cfg(target_os = "macos")]
140    type SoftmaxAcc = f16;
141    #[cfg(target_os = "macos")]
142    type Accumulator = f16;
143    #[cfg(not(target_os = "macos"))]
144    type SoftmaxAcc = f32;
145    #[cfg(not(target_os = "macos"))]
146    type Accumulator = f32;
147    type Mask = u8;
148    type MaskSize = Const<0>;
149    type Out = f16;
150}
151
152impl AttentionPrecision for flex32 {
153    type Query = flex32;
154    type Key = flex32;
155    type Value = flex32;
156    type KVTile = f16;
157    type SoftmaxLhs = f16;
158    #[cfg(target_os = "macos")]
159    type SoftmaxAcc = f16;
160    #[cfg(target_os = "macos")]
161    type Accumulator = f16;
162    #[cfg(not(target_os = "macos"))]
163    type SoftmaxAcc = f32;
164    #[cfg(not(target_os = "macos"))]
165    type Accumulator = f32;
166    type Mask = u8;
167    type MaskSize = Const<0>;
168    type Out = f32;
169}
170
171impl AttentionPrecision for bf16 {
172    type Query = bf16;
173    type Key = bf16;
174    type Value = bf16;
175    type KVTile = bf16;
176    type SoftmaxLhs = bf16;
177    #[cfg(target_os = "macos")]
178    type SoftmaxAcc = bf16;
179    #[cfg(target_os = "macos")]
180    type Accumulator = bf16;
181    #[cfg(not(target_os = "macos"))]
182    type SoftmaxAcc = f32;
183    #[cfg(not(target_os = "macos"))]
184    type Accumulator = f32;
185    type Mask = u8;
186    type MaskSize = Const<0>;
187    type Out = bf16;
188}
189
190impl AttentionPrecision for f32 {
191    type Query = f32;
192    type Key = f32;
193    type Value = f32;
194    type KVTile = f32;
195    type SoftmaxAcc = f32;
196    type SoftmaxLhs = f32;
197    type Accumulator = f32;
198    type Mask = u8;
199    type MaskSize = Const<0>;
200    type Out = f32;
201}
202
203impl AttentionPrecision for f64 {
204    type Query = f64;
205    type Key = f64;
206    type Value = f64;
207    type KVTile = f32;
208    type SoftmaxAcc = f32;
209    type SoftmaxLhs = f32;
210    type Accumulator = f32;
211    type Mask = u8;
212    type MaskSize = Const<0>;
213    type Out = f64;
214}
215
216impl<
217    QG: Float,
218    QGS: Size,
219    QT: Float,
220    KG: Float,
221    KGS: Size,
222    KS: Float,
223    KSS: Size,
224    VG: Float,
225    VGS: Size,
226    VS: Float,
227    VSS: Size,
228    KVT: Float,
229    SM: Float,
230    SML: Float,
231    ACC: Float,
232    MSK: Numeric,
233    MSKS: Size,
234    OG: Float,
235    OGS: Size,
236    OS: Float,
237    OSS: Size,
238> AttentionPrecision
239    for (
240        (QG, QGS, QT),
241        (KG, KGS, KS, KSS),
242        (VG, VGS, VS, VSS),
243        KVT,
244        SM,
245        SML,
246        ACC,
247        MSK,
248        MSKS,
249        (OG, OGS, OS, OSS),
250    )
251{
252    type Query = (QG, QGS, QT);
253    type Key = (KG, KGS, KS, KSS);
254    type Value = (VG, VGS, VS, VSS);
255    type KVTile = KVT;
256    type SoftmaxAcc = SM;
257    type SoftmaxLhs = SML;
258    type Accumulator = ACC;
259    type Mask = MSK;
260    type MaskSize = MSKS;
261    type Out = (OG, OGS, OS, OSS);
262}
263
264pub mod launch_types {
265    use super::*;
266
267    define_scalar!(pub QG);
268    define_scalar!(pub KG);
269    define_scalar!(pub VG);
270    define_scalar!(pub MSK);
271    define_scalar!(pub OG);
272
273    define_size!(pub QGS);
274    define_size!(pub KGS);
275    define_size!(pub VGS);
276    define_size!(pub MSKS);
277    define_size!(pub OGS);
278
279    /// Input argument
280    pub type InputArg<AA> =
281        <AA as AttentionArgs>::Input<(QG, QGS), (KG, KGS), (VG, VGS), (MSK, MSKS)>;
282
283    /// Output argument
284    pub type OutputArg<AA> = <AA as AttentionArgs>::Output<(OG, OGS)>;
285}
286
287pub use launch_types::{InputArg, OutputArg};
288
289/// Input runtime argument
290pub type InputRuntimeArg<AA, R> = <InputArg<AA> as LaunchArg>::RuntimeArg<R>;
291
292/// Output runtime argument
293pub type OutputRuntimeArg<AA, R> = <OutputArg<AA> as LaunchArg>::RuntimeArg<R>;
294
295pub mod attention_types {
296    use ruda_kernel::dsl::prelude::*;
297
298    use crate::attention::kernel_ir::definition::{
299        AttentionPrecision, AttentionSpec, QueryPrecision, StagedMatrixPrecision,
300    };
301
302    // ==================== Per-operand precision grouping ====================
303
304    pub type Query<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::Query;
305    pub type Key<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::Key;
306    pub type Value<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::Value;
307    pub type Out<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::Out;
308
309    // ==================== QUERY ====================
310
311    // Element / Size splits
312    pub type QG<AS> = <Query<AS> as QueryPrecision>::Global;
313    pub type QGS<AS> = <Query<AS> as QueryPrecision>::GlobalSize;
314    pub type QT<AS> = <Query<AS> as QueryPrecision>::Tile;
315
316    // Vector form
317    pub type QGV<AS> = Vector<QG<AS>, QGS<AS>>;
318
319    // ==================== KEY ====================
320
321    // Element / Size splits
322    pub type KG<AS> = <Key<AS> as StagedMatrixPrecision>::Global;
323    pub type KGS<AS> = <Key<AS> as StagedMatrixPrecision>::GlobalSize;
324    pub type KS<AS> = <Key<AS> as StagedMatrixPrecision>::Stage;
325    pub type KSS<AS> = <Key<AS> as StagedMatrixPrecision>::StageSize;
326
327    // Vector forms
328    pub type KGV<AS> = Vector<KG<AS>, KGS<AS>>;
329    pub type KSV<AS> = Vector<KS<AS>, KSS<AS>>;
330
331    // ==================== VALUE ====================
332
333    // Element / Size splits
334    pub type VG<AS> = <Value<AS> as StagedMatrixPrecision>::Global;
335    pub type VGS<AS> = <Value<AS> as StagedMatrixPrecision>::GlobalSize;
336    pub type VS<AS> = <Value<AS> as StagedMatrixPrecision>::Stage;
337    pub type VSS<AS> = <Value<AS> as StagedMatrixPrecision>::StageSize;
338
339    // Vector forms
340    pub type VGV<AS> = Vector<VG<AS>, VGS<AS>>;
341    pub type VSV<AS> = Vector<VS<AS>, VSS<AS>>;
342
343    // ==================== KV TILE / SOFTMAX / ACCUMULATOR ====================
344
345    pub type KVT<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::KVTile;
346    pub type SM<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::SoftmaxAcc;
347    pub type SML<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::SoftmaxLhs;
348    pub type ACC<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::Accumulator;
349
350    // ==================== MASK ====================
351
352    pub type MSK<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::Mask;
353    pub type MSKS<AS> = <<AS as AttentionSpec>::Precision as AttentionPrecision>::MaskSize;
354
355    // Vector form
356    pub type MSKV<AS> = Vector<MSK<AS>, MSKS<AS>>;
357
358    // ==================== OUT ====================
359
360    // Element / Size splits
361    pub type OG<AS> = <Out<AS> as StagedMatrixPrecision>::Global;
362    pub type OGS<AS> = <Out<AS> as StagedMatrixPrecision>::GlobalSize;
363    pub type OS<AS> = <Out<AS> as StagedMatrixPrecision>::Stage;
364    pub type OSS<AS> = <Out<AS> as StagedMatrixPrecision>::StageSize;
365
366    // Vector forms
367    pub type OGV<AS> = Vector<OG<AS>, OGS<AS>>;
368    pub type OSV<AS> = Vector<OS<AS>, OSS<AS>>;
369}
370
371pub type Args<MS> = <MS as AttentionSpec>::Args;
372
373#[derive(Debug, Clone, Eq, PartialEq, Hash)]
374pub struct AttentionElems {
375    pub query_global: StorageType,
376    pub query_tile: StorageType,
377    pub key_global: StorageType,
378    pub key_stage: StorageType,
379    pub value_global: StorageType,
380    pub value_stage: StorageType,
381    pub key_value_tile: StorageType,
382    pub softmax_acc: StorageType,
383    pub softmax_lhs: StorageType,
384    pub accumulator: StorageType,
385    pub mask: StorageType,
386    pub out_global: StorageType,
387    pub out_stage: StorageType,
388}
389
390impl AttentionElems {
391    pub fn from_global_types(
392        global_dtypes: &AttentionGlobalTypes,
393        tile_type: StorageType,
394        accumulator_precision: &AccumulatorPrecision,
395    ) -> AttentionElems {
396        let accumulator = match accumulator_precision {
397            AccumulatorPrecision::Strict(storage_type) => *storage_type,
398            AccumulatorPrecision::Loose => AccumulatorPrecision::default_accumulator_type(),
399        };
400
401        Self {
402            query_global: global_dtypes.query,
403            query_tile: tile_type,
404            key_global: global_dtypes.key,
405            key_stage: tile_type,
406            value_global: global_dtypes.value,
407            value_stage: tile_type,
408            key_value_tile: tile_type,
409            softmax_acc: accumulator,
410            softmax_lhs: tile_type,
411            accumulator,
412            mask: global_dtypes.mask,
413            out_global: global_dtypes.out,
414            out_stage: global_dtypes.out,
415        }
416    }
417}
418
419impl From<&AttentionElems> for [StorageType; 5] {
420    fn from(elems: &AttentionElems) -> Self {
421        [
422            elems.query_global,
423            elems.key_global,
424            elems.value_global,
425            elems.mask,
426            elems.out_global,
427        ]
428    }
429}
rudnn/attention/kernel_ir/definition/spec.rs

rudnn/attention/kernel_ir/definition/
spec.rs
