kyro 0.1.1

A high-performance ML inference engine
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

#![allow(dead_code)]

pub struct PipelineContext {
    pub rank: usize,
    pub world_size: usize,
    pub start_layer: usize,
    pub end_layer: usize,
}

impl PipelineContext {
    pub fn new(rank: usize, world_size: usize, total_layers: usize) -> Self {
        let layers_per_gpu = total_layers.div_ceil(world_size);
        let start_layer = rank * layers_per_gpu;
        let end_layer = std::cmp::min(start_layer + layers_per_gpu, total_layers);

        Self {
            rank,
            world_size,
            start_layer,
            end_layer,
        }
    }

    pub fn is_first_stage(&self) -> bool {
        self.rank == 0
    }

    pub fn is_last_stage(&self) -> bool {
        self.rank == self.world_size - 1
    }

    pub fn should_process_layer(&self, index: usize) -> bool {
        index >= self.start_layer && index < self.end_layer
    }
}