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

use mircuda::{DeviceBuffer, LaunchConfig, Stream, bf16};

use super::{ClampedRoutedAttention, narrow};
use crate::{PagedPrefillBatch, Result};

impl ClampedRoutedAttention {
    #[allow(clippy::too_many_arguments)]
    pub(super) fn execute_fmha(
        &self,
        stream: &Stream,
        query: &DeviceBuffer<bf16>,
        key_pages: &DeviceBuffer<u8>,
        value_pages: &DeviceBuffer<u8>,
        batch: &PagedPrefillBatch,
        tables: &DeviceBuffer<u32>,
        sinks: &DeviceBuffer<bf16>,
        softmax_lse: &mut DeviceBuffer<bf16>,
        output: &mut DeviceBuffer<bf16>,
        scale: f32,
    ) -> Result<bool> {
        let Some(fmha) = &self.fmha else {
            return Ok(false);
        };
        let max_context_tokens = batch
            .max_context_tokens()
            .div_ceil(128)
            .saturating_mul(128)
            .min(batch.max_blocks().saturating_mul(batch.cache_config().block_size));
        fmha.execute_paged_varlen(
            stream,
            query,
            key_pages,
            value_pages,
            output,
            batch.query_starts(),
            batch.token_counts(),
            batch.context_starts(),
            tables,
            softmax_lse,
            batch.active(),
            batch.tokens(),
            batch.max_query_tokens(),
            max_context_tokens,
            batch.max_blocks(),
            batch.cache_config().block_size,
            scale,
        )?;
        self.sink_scale.launch(
            stream,
            LaunchConfig::for_elements(batch.tokens() * self.query_heads * self.head_dim, 256)?,
            (
                output,
                &*softmax_lse,
                sinks,
                narrow(batch.tokens())?,
                narrow(self.query_heads)?,
                narrow(self.head_dim)?,
            ),
        )?;
        Ok(true)
    }
}