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
36
37
38
39

#![allow(dead_code)]

use candle_core::{Device, Result, Tensor};

pub struct PagedAttention {
    pub block_size: usize,
    pub num_heads: usize,
    pub head_dim: usize,
}

impl PagedAttention {
    pub fn new(block_size: usize, num_heads: usize, head_dim: usize) -> Self {
        Self {
            block_size,
            num_heads,
            head_dim,
        }
    }

    pub fn forward(
        &self,
        query: &Tensor,
        _key_cache: &Tensor,
        _value_cache: &Tensor,
        _block_table: &Tensor,
        _context_lens: &Tensor,
    ) -> Result<Tensor> {
        let device = query.device();

        match device {
            Device::Cuda(_) => self.software_paged_attention(query),
            _ => self.software_paged_attention(query),
        }
    }

    fn software_paged_attention(&self, query: &Tensor) -> Result<Tensor> {
        Ok(query.clone())
    }
}