libmir-cuda 0.3.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

use mircuda::{DeviceBuffer, bf16};
use runtime::kv::BlockTable;

use super::ClampedRoutedLayerExecution;
use crate::{PagedKvCache, Result, kernels::ClampedRoutedSplitDecode};

impl ClampedRoutedLayerExecution {
    #[allow(clippy::too_many_arguments)]
    pub(super) fn attend(
        &self,
        cache: &PagedKvCache,
        current_keys: &DeviceBuffer<bf16>,
        current_values: &DeviceBuffer<bf16>,
        table: &BlockTable,
        table_device: &DeviceBuffer<u32>,
        sinks: &DeviceBuffer<bf16>,
        start: usize,
        query: &DeviceBuffer<bf16>,
        attended: &mut DeviceBuffer<bf16>,
        split_decode: &mut Option<ClampedRoutedSplitDecode>,
    ) -> Result<()> {
        let blocks = table.blocks().len();
        if self.tokens == 1 {
            if !cache.is_windowed()
                && let Some(split) = split_decode.as_mut()
            {
                split.ensure_tuned(
                    &self.attention,
                    &self.stream,
                    query,
                    cache.key_pages(),
                    cache.value_pages(),
                    table_device,
                    sinks,
                    attended,
                    table.token_len(),
                    blocks,
                    self.window,
                    self.config.scale,
                );
                if split.execute(
                    &self.stream,
                    query,
                    cache.key_pages(),
                    cache.value_pages(),
                    table_device,
                    sinks,
                    attended,
                    table.token_len(),
                    blocks,
                    self.window,
                    self.config.scale,
                )? {
                    return Ok(());
                }
            }
            self.attention.execute(
                &self.stream,
                query,
                current_keys,
                current_values,
                cache.key_pages(),
                cache.value_pages(),
                table_device,
                sinks,
                attended,
                table.token_len(),
                blocks,
                self.window,
                self.config.scale,
            )
        } else {
            self.attention.execute_prefill(
                &self.stream,
                query,
                current_keys,
                current_values,
                cache.key_pages(),
                cache.value_pages(),
                table_device,
                sinks,
                attended,
                self.tokens,
                start,
                blocks,
                self.window,
                self.config.scale,
            )
        }
    }
}