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

use std::time::Duration;

use runtime::tuning::materially_faster;

use crate::AttentionExecution;

pub struct SplitMeasurement {
    pub partition_tokens: usize,
    pub timings: Vec<Duration>,
    pub score: Duration,
}

pub fn select_execution(
    fallback: AttentionExecution,
    fallback_partition: usize,
    max_context_tokens: usize,
    contexts: &[usize],
    direct: &[Duration],
    splits: &[SplitMeasurement],
    minimum_bps: u16,
) -> AttentionExecution {
    let Some(fastest) = splits.iter().min_by_key(|candidate| candidate.score) else {
        return fallback;
    };
    let baseline = splits
        .iter()
        .find(|candidate| candidate.partition_tokens == fallback_partition)
        .unwrap_or(fastest);
    let selected = if materially_faster(fastest.score, baseline.score, minimum_bps) {
        fastest
    } else {
        baseline
    };
    let crossover = contexts.iter().copied().zip(direct).zip(&selected.timings).find_map(
        |((tokens, direct), split)| {
            materially_faster(*split, *direct, minimum_bps).then_some(tokens)
        },
    );
    match fallback {
        AttentionExecution::Direct => {
            crossover.map_or(AttentionExecution::Direct, |threshold| AttentionExecution::SplitKv {
                partition_tokens: selected.partition_tokens,
                threshold_tokens: threshold,
            })
        },
        AttentionExecution::SplitKv { threshold_tokens, .. } => {
            let observed_limit = contexts.last().copied().unwrap_or(0);
            let threshold_tokens = crossover.unwrap_or_else(|| {
                observed_limit
                    .saturating_add(1)
                    .min(max_context_tokens.saturating_add(1))
                    .max(threshold_tokens)
            });
            AttentionExecution::SplitKv {
                partition_tokens: selected.partition_tokens,
                threshold_tokens,
            }
        },
    }
}

pub fn execution_average(
    execution: AttentionExecution,
    direct: &[Duration],
    splits: &[SplitMeasurement],
) -> Duration {
    let timings = match execution {
        AttentionExecution::Direct => direct,
        AttentionExecution::SplitKv { partition_tokens, .. } => splits
            .iter()
            .find(|candidate| candidate.partition_tokens == partition_tokens)
            .map_or(direct, |candidate| candidate.timings.as_slice()),
    };
    let count = u32::try_from(timings.len()).unwrap_or(u32::MAX).max(1);
    timings.iter().copied().sum::<Duration>() / count
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::{SplitMeasurement, select_execution};
    use crate::AttentionExecution;

    fn split(partition_tokens: usize, micros: &[u64]) -> SplitMeasurement {
        let timings = micros.iter().map(|value| Duration::from_micros(*value)).collect::<Vec<_>>();
        let score = timings.iter().copied().sum();
        SplitMeasurement { partition_tokens, timings, score }
    }

    #[test]
    fn learns_partition_and_first_material_crossover() {
        let selected = select_execution(
            AttentionExecution::SplitKv {
                partition_tokens: 256,
                threshold_tokens: 65,
            },
            256,
            4_096,
            &[64, 256, 1_024],
            &[Duration::from_micros(10), Duration::from_micros(30), Duration::from_micros(100)],
            &[split(128, &[20, 20, 40]), split(256, &[25, 28, 70])],
            300,
        );
        assert_eq!(
            selected,
            AttentionExecution::SplitKv {
                partition_tokens: 128,
                threshold_tokens: 256
            }
        );
    }

    #[test]
    fn direct_fallback_is_retained_without_a_crossover() {
        let selected = select_execution(
            AttentionExecution::Direct,
            256,
            512,
            &[64, 256],
            &[Duration::from_micros(10), Duration::from_micros(20)],
            &[split(256, &[20, 30])],
            300,
        );
        assert_eq!(selected, AttentionExecution::Direct);
    }
}