lift-tensor 0.2.0

LIFT-TENSOR: AI dialect — tensor types, attention, KV cache, MoE, quantisation
Documentation 
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
94
95
96
97
98
99
100
101

use serde::{Serialize, Deserialize};
use lift_core::types::{DataType, Dimension, MemoryLayout};

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum TensorType {
    Tensor {
        shape: Vec<Dimension>,
        dtype: DataType,
        layout: MemoryLayout,
    },
    AttentionTensor {
        batch: Dimension,
        seq_len: Dimension,
        num_heads: usize,
        head_dim: usize,
        dtype: DataType,
    },
    KVCache {
        max_seq: Dimension,
        num_heads: usize,
        head_dim: usize,
        dtype: DataType,
        is_paged: bool,
    },
    SparseTensor {
        num_experts: usize,
        capacity: usize,
        dtype: DataType,
    },
}

impl TensorType {
    pub fn dtype(&self) -> DataType {
        match self {
            TensorType::Tensor { dtype, .. } => *dtype,
            TensorType::AttentionTensor { dtype, .. } => *dtype,
            TensorType::KVCache { dtype, .. } => *dtype,
            TensorType::SparseTensor { dtype, .. } => *dtype,
        }
    }

    pub fn element_count(&self) -> Option<usize> {
        match self {
            TensorType::Tensor { shape, .. } => {
                let mut count = 1usize;
                for dim in shape {
                    match dim.static_value() {
                        Some(v) => count = count.checked_mul(v)?,
                        None => return None,
                    }
                }
                Some(count)
            }
            TensorType::AttentionTensor { batch, seq_len, num_heads, head_dim, .. } => {
                let b = batch.static_value()?;
                let s = seq_len.static_value()?;
                Some(b * s * num_heads * head_dim)
            }
            _ => None,
        }
    }

    pub fn size_bytes(&self) -> Option<usize> {
        Some(self.element_count()? * self.dtype().byte_size())
    }

    pub fn rank(&self) -> Option<usize> {
        match self {
            TensorType::Tensor { shape, .. } => Some(shape.len()),
            TensorType::AttentionTensor { .. } => Some(4), // [batch, seq, heads, dim]
            TensorType::KVCache { .. } => Some(4),
            TensorType::SparseTensor { .. } => Some(3),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AttentionImpl {
    Standard,
    FlashAttentionV2,
    FlashAttentionV3,
    PagedAttention,
    SDPA,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum QuantizationMode {
    DynamicINT8,
    StaticINT8,
    FP8E4M3,
    FP8E5M2,
    INT4,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ParallelismStrategy {
    DataParallel,
    TensorParallel,
    PipelineParallel,
    SequenceParallel,
}