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

use std::{collections::HashMap, time::Duration};

use super::{
    AttentionProfileRequest, AttentionRuntimeEntry, CudaAutoTuner, storage::StoredAttentionEntry,
};
use crate::{AttentionExecution, PlanSource};

impl CudaAutoTuner {
    pub(crate) fn lookup_attention(
        &self,
        request: AttentionProfileRequest,
    ) -> Option<(AttentionExecution, PlanSource)> {
        if self.inner.config.mode == super::CudaTuningMode::Disabled {
            return None;
        }
        self.inner
            .state
            .lock()
            .ok()?
            .attention
            .get(&request)
            .map(|entry| (entry.execution, entry.source))
    }

    pub(crate) fn claim_attention(&self, request: AttentionProfileRequest) -> bool {
        let Ok(mut state) = self.inner.state.lock() else {
            return false;
        };
        self.inner.config.mode == super::CudaTuningMode::Startup
            && !state.sealed
            && state.budget.available()
            && !state.attention.contains_key(&request)
            && state.attention_inflight.insert(request)
    }

    pub(crate) fn record_attention(
        &self,
        request: AttentionProfileRequest,
        execution: AttentionExecution,
        average: Duration,
        tuning_elapsed: Duration,
    ) {
        let snapshot = {
            let Ok(mut state) = self.inner.state.lock() else {
                return;
            };
            state.attention_inflight.remove(&request);
            state.budget.consume(tuning_elapsed);
            state.attention.insert(
                request,
                AttentionRuntimeEntry {
                    execution,
                    source: PlanSource::MeasuredStartup,
                    average_ns: u64::try_from(average.as_nanos()).unwrap_or(u64::MAX),
                },
            );
            Self::snapshot(&state)
        };
        self.persist(snapshot);
    }

    pub(crate) fn abandon_attention(&self, request: AttentionProfileRequest) {
        if let Ok(mut state) = self.inner.state.lock() {
            state.attention_inflight.remove(&request);
        }
    }
}

pub(super) fn stored_entries(
    entries: &HashMap<AttentionProfileRequest, AttentionRuntimeEntry>,
) -> Vec<StoredAttentionEntry> {
    let mut stored = entries
        .iter()
        .map(|(request, entry)| StoredAttentionEntry {
            request: *request,
            execution: entry.execution,
            average_ns: entry.average_ns,
        })
        .collect::<Vec<_>>();
    stored.sort_by_key(|entry| {
        (
            entry.request.family as u8,
            entry.request.plan.query_heads,
            entry.request.plan.kv_heads,
            entry.request.plan.head_dim,
            entry.request.plan.value_head_dim,
            entry.request.plan.max_context_tokens,
            entry.request.block_size,
            entry.request.window_tokens,
        )
    });
    stored
}