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

use super::ClampedRoutedExecutionPlan;
use crate::PagedPrefillBatch;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(in crate::backend::clamped_routed) enum ClampedRoutedDecodeSignature {
    Direct { context_bucket: usize },
    Split { partitions: usize },
}

impl ClampedRoutedExecutionPlan {
    pub(in crate::backend::clamped_routed) fn decode_signature(
        &self,
        batch: &PagedPrefillBatch,
    ) -> ClampedRoutedDecodeSignature {
        let partitions = self
            .batch_split_decode
            .as_ref()
            .map_or(0, |split| split.capture_partitions(batch));
        if partitions > 0 {
            ClampedRoutedDecodeSignature::Split { partitions }
        } else {
            ClampedRoutedDecodeSignature::Direct {
                context_bucket: batch.max_context_tokens().div_ceil(128),
            }
        }
    }
}