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

use std::time::{Duration, Instant};

use runtime::backend::{PrefillOutput, PrefillRequest};

use super::{Sequence, prefix::PrefixReuse};
use crate::{Error, Result};

impl Sequence {
    pub(super) fn new(request: PrefillRequest, prefix: PrefixReuse, runner_wait: Duration) -> Self {
        let step_table = request.block_table.clone();
        Self {
            request,
            consumed: prefix.tokens,
            chunks: 0,
            prefix_tokens: prefix.tokens,
            checkpoint_restored: prefix.checkpoint_restored,
            runner_wait,
            completed_at: None,
            step_table,
            output: None,
        }
    }

    pub(super) fn pending(&self) -> bool {
        self.consumed < self.request.prompt_tokens.len()
    }

    pub(super) fn finish(
        mut self,
        loaded: &crate::engine::model::LoadedModel,
        started: Instant,
    ) -> Result<PrefillOutput> {
        let elapsed = started.elapsed();
        let output = self
            .output
            .take()
            .ok_or(Error::InvalidDecoderKernel("CUDA prefill produced no output"))?;
        let completion = self
            .completed_at
            .map_or(elapsed, |completed| completed.saturating_duration_since(started));
        loaded.register_session(self.request.session_id)?;
        tracing::debug!(
            backend = "cuda",
            session = %self.request.session_id,
            prompt_tokens = self.request.prompt_tokens.len(),
            prefix_cache_tokens = self.prefix_tokens,
            checkpoint_restored = self.checkpoint_restored,
            runner_wait_ms = self.runner_wait.as_secs_f64() * 1_000.0,
            completion_ms = completion.as_secs_f64() * 1_000.0,
            cohort_elapsed_ms = elapsed.as_secs_f64() * 1_000.0,
            chunks = self.chunks,
            "completed CUDA prefill request"
        );
        Ok(PrefillOutput {
            accepted_tokens: self.request.prompt_tokens.len(),
            next_token: output.token,
            trace: Some(format!(
                "cuda.prefill=token-budget-{};prefix_cache_tokens={};checkpoint_restored={}",
                if self.checkpoint_restored {
                    "checkpoint-completion-first"
                } else {
                    "round-robin"
                },
                self.prefix_tokens,
                self.checkpoint_restored
            )),
            logits: output.logits,
            candidates: None,
        })
    }
}