libmir-cuda 0.1.0

CUDA inference backend for libmir
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

use super::{AttentionArguments, KvArguments, QkvPostprocessArguments, Resources};
use crate::{
    Result,
    kernels::{MergeAttentionArguments, SplitAttentionArguments},
};

pub(super) fn qkv_arguments(resources: &mut Resources) -> QkvPostprocessArguments<'_> {
    let geometry = resources.geometry;
    let dynamic = resources.dynamic;
    let weights = resources.weights.borrow();
    let scratch = &mut resources.attention.scratch;
    let (query, key, value) = if geometry.separate_qkv {
        (&scratch.qkv_separate[0], &scratch.qkv_separate[1], &scratch.qkv_separate[2])
    } else {
        (&scratch.qkv, &scratch.qkv, &scratch.qkv)
    };
    (
        query,
        key,
        value,
        weights.query_norm,
        weights.key_norm,
        &mut scratch.query_rope,
        &mut scratch.key_rope,
        &mut scratch.value_norm,
        1,
        geometry.query_heads,
        geometry.kv_heads,
        geometry.head_dim,
        geometry.value_head_dim,
        geometry.rotary_dim,
        geometry.pairing_dim,
        dynamic.position,
        geometry.theta,
        geometry.epsilon,
        u32::from(geometry.separate_qkv),
        u32::from(geometry.normalization.query),
        u32::from(geometry.normalization.key),
        u32::from(geometry.normalization.value),
    )
}

pub(super) fn kv_arguments(resources: &mut Resources) -> KvArguments<'_> {
    let geometry = resources.geometry;
    let dynamic = resources.dynamic;
    let attention = &mut resources.attention;
    let (key_pages, value_pages) = attention.cache.pages_mut();
    (
        &attention.scratch.key_rope,
        &attention.scratch.value_norm,
        key_pages,
        value_pages,
        dynamic.local_start,
        1,
        dynamic.physical_block,
        dynamic.page_start,
        geometry.block_size_abi,
        geometry.kv_heads,
        geometry.head_dim,
        geometry.value_head_dim,
    )
}

pub(super) fn attention_arguments(resources: &mut Resources) -> AttentionArguments<'_> {
    let geometry = resources.geometry;
    let dynamic = resources.dynamic;
    let attention = &mut resources.attention;
    (
        &attention.scratch.query_rope,
        attention.cache.key_pages(),
        attention.cache.value_pages(),
        attention.attention.table_device(),
        &mut attention.scratch.attention,
        dynamic.token_count,
        dynamic.block_count,
        geometry.block_size_abi,
        geometry.query_heads,
        geometry.kv_heads,
        geometry.head_dim,
        geometry.value_head_dim,
        geometry.window,
        geometry.scale,
        geometry.split_threshold,
    )
}

pub(super) fn split_attention_arguments(
    resources: &mut Resources,
) -> Result<SplitAttentionArguments<'_>> {
    let geometry = resources.geometry;
    let dynamic = resources.dynamic;
    let attention = &mut resources.attention;
    attention.attention.captured_split_arguments(
        &attention.scratch.query_rope,
        &attention.cache,
        dynamic.token_count,
        dynamic.block_count,
        geometry.window,
        geometry.scale,
    )
}

pub(super) fn merge_attention_arguments(
    resources: &mut Resources,
) -> Result<MergeAttentionArguments<'_>> {
    let geometry = resources.geometry;
    let dynamic = resources.dynamic;
    let attention = &mut resources.attention;
    attention.attention.captured_merge_arguments(
        &mut attention.scratch.attention,
        dynamic.token_count,
        geometry.window,
    )
}