runtime/
types.rs

1//! Core types for UniLLM runtime
2//!
3//! This module defines fundamental data structures used throughout the runtime.
4
5use std::collections::HashMap;
6use std::fmt;
7
8
9/// Data types supported by tensors
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum DataType {
12    Float32,
13    Float16,
14    BFloat16,
15    Int32,
16    Int64,
17    Int8,
18    Int4,
19    Bool,
20}
21
22impl DataType {
23    pub fn size_bytes(&self) -> usize {
24        match self {
25            DataType::Float32 | DataType::Int32 => 4,
26            DataType::Float16 | DataType::BFloat16 => 2,
27            DataType::Int64 => 8,
28            DataType::Int8 => 1,
29            DataType::Int4 => 1, // Packed
30            DataType::Bool => 1,
31        }
32    }
33}
34
35/// Device types for tensor placement
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub enum Device {
38    CPU,
39    CUDA(usize),  // GPU ID
40    ROCm(usize),  // GPU ID
41    Intel(usize), // XPU ID
42    Metal(usize), // GPU ID
43}
44
45impl fmt::Display for Device {
46    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47        match self {
48            Device::CPU => write!(f, "cpu"),
49            Device::CUDA(id) => write!(f, "cuda:{}", id),
50            Device::ROCm(id) => write!(f, "rocm:{}", id),
51            Device::Intel(id) => write!(f, "intel:{}", id),
52            Device::Metal(id) => write!(f, "metal:{}", id),
53        }
54    }
55}
56
57/// Tensor structure with real data storage
58#[derive(Debug, Clone)]
59pub struct Tensor {
60    pub shape: Vec<usize>,
61    pub dtype: DataType,
62    pub device: Device,
63    pub data: TensorData,
64    pub strides: Vec<usize>,
65}
66
67/// Tensor data storage
68#[derive(Debug, Clone)]
69pub enum TensorData {
70    F32(Vec<f32>),
71    F16(Vec<u16>), // f16 as u16
72    BF16(Vec<u16>), // bf16 as u16
73    I32(Vec<i32>),
74    I64(Vec<i64>),
75    I8(Vec<i8>),
76    Bool(Vec<bool>),
77    Unallocated, // For lazy loading
78}
79
80impl Tensor {
81    pub fn new(shape: Vec<usize>, dtype: DataType, device: Device) -> Self {
82        let strides = calculate_strides(&shape);
83        let numel = shape.iter().product::<usize>();
84
85        let data = match dtype {
86            DataType::Float32 => TensorData::F32(vec![0.0; numel]),
87            DataType::Float16 => TensorData::F16(vec![0; numel]),
88            DataType::BFloat16 => TensorData::BF16(vec![0; numel]),
89            DataType::Int32 => TensorData::I32(vec![0; numel]),
90            DataType::Int64 => TensorData::I64(vec![0; numel]),
91            DataType::Int8 => TensorData::I8(vec![0; numel]),
92            DataType::Bool => TensorData::Bool(vec![false; numel]),
93            DataType::Int4 => TensorData::I8(vec![0; (numel + 1) / 2]), // Packed
94        };
95
96        Self {
97            shape,
98            dtype,
99            device,
100            data,
101            strides,
102        }
103    }
104
105    pub fn zeros(shape: &[usize]) -> Self {
106        Self::new(shape.to_vec(), DataType::Float32, Device::CPU)
107    }
108
109    pub fn ones(shape: &[usize]) -> Self {
110        let mut tensor = Self::new(shape.to_vec(), DataType::Float32, Device::CPU);
111        if let TensorData::F32(ref mut data) = tensor.data {
112            data.fill(1.0);
113        }
114        tensor
115    }
116
117    pub fn from_data(shape: Vec<usize>, data: TensorData, device: Device) -> Self {
118        let dtype = match &data {
119            TensorData::F32(_) => DataType::Float32,
120            TensorData::F16(_) => DataType::Float16,
121            TensorData::BF16(_) => DataType::BFloat16,
122            TensorData::I32(_) => DataType::Int32,
123            TensorData::I64(_) => DataType::Int64,
124            TensorData::I8(_) => DataType::Int8,
125            TensorData::Bool(_) => DataType::Bool,
126            TensorData::Unallocated => DataType::Float32, // Default
127        };
128
129        let strides = calculate_strides(&shape);
130
131        Self {
132            shape,
133            dtype,
134            device,
135            data,
136            strides,
137        }
138    }
139
140    pub fn numel(&self) -> usize {
141        self.shape.iter().product()
142    }
143
144    pub fn size_bytes(&self) -> usize {
145        self.numel() * self.dtype.size_bytes()
146    }
147
148    pub fn index_select(&self, _dim: usize, _indices: &Tensor) -> anyhow::Result<Tensor> {
149        // Placeholder implementation
150        Ok(self.clone())
151    }
152
153    pub fn matmul(&self, other: &Tensor) -> anyhow::Result<Tensor> {
154        // Basic matrix multiplication for F32 tensors
155        if self.shape.len() != 2 || other.shape.len() != 2 {
156            return Err(anyhow::anyhow!("matmul requires 2D tensors"));
157        }
158
159        if self.shape[1] != other.shape[0] {
160            return Err(anyhow::anyhow!("matmul dimension mismatch: {} x {} != {} x {}",
161                self.shape[0], self.shape[1], other.shape[0], other.shape[1]));
162        }
163
164        let result_shape = vec![self.shape[0], other.shape[1]];
165        let mut result = Tensor::zeros(&result_shape);
166
167        // Perform computation based on data types
168        match (&self.data, &other.data, &mut result.data) {
169            (TensorData::F32(a), TensorData::F32(b), TensorData::F32(c)) => {
170                let m = self.shape[0];
171                let n = other.shape[1];
172                let k = self.shape[1];
173
174                for i in 0..m {
175                    for j in 0..n {
176                        let mut sum = 0.0;
177                        for kk in 0..k {
178                            sum += a[i * k + kk] * b[kk * n + j];
179                        }
180                        c[i * n + j] = sum;
181                    }
182                }
183            }
184            _ => return Err(anyhow::anyhow!("matmul not implemented for these data types")),
185        }
186
187        Ok(result)
188    }
189
190    pub fn add(&self, other: &Tensor) -> anyhow::Result<Tensor> {
191        if self.shape != other.shape {
192            return Err(anyhow::anyhow!("add requires tensors with same shape"));
193        }
194
195        let mut result = self.clone();
196
197        match (&self.data, &other.data, &mut result.data) {
198            (TensorData::F32(a), TensorData::F32(b), TensorData::F32(c)) => {
199                for i in 0..a.len() {
200                    c[i] = a[i] + b[i];
201                }
202            }
203            _ => return Err(anyhow::anyhow!("add not implemented for these data types")),
204        }
205
206        Ok(result)
207    }
208
209    pub fn arange(start: i64, end: i64) -> Self {
210        let size = (end - start) as usize;
211        Self::new(vec![size], DataType::Int64, Device::CPU)
212    }
213
214    pub fn layer_norm(&self, _weight: &Tensor) -> anyhow::Result<Tensor> {
215        // Placeholder implementation
216        Ok(self.clone())
217    }
218
219    pub fn gelu(&self) -> anyhow::Result<Tensor> {
220        // Placeholder implementation
221        Ok(self.clone())
222    }
223
224    pub fn tanh(&self) -> anyhow::Result<Tensor> {
225        // Placeholder implementation
226        Ok(self.clone())
227    }
228}
229
230/// Model configuration
231#[derive(Debug, Clone)]
232pub struct ModelConfig {
233    pub model_name: String,
234    pub model_path: String,
235    pub max_sequence_length: usize,
236    pub vocabulary_size: usize,
237    pub num_layers: usize,
238    pub num_heads: usize,
239    pub head_dim: usize,
240    pub hidden_size: usize,
241    pub intermediate_size: usize,
242    pub dtype: DataType,
243}
244
245/// Model features that can be supported
246#[derive(Debug, Clone, Copy, PartialEq, Eq)]
247pub enum ModelFeature {
248    FlashAttention,
249    GroupedQueryAttention,
250    RotaryEmbedding,
251    RMSNorm,
252    LayerNorm,
253    SwiGLU,
254    GELU,
255    PrefixCaching,
256    ChunkedPrefill,
257    DynamicBatching,
258    ContinuousBatching,
259    LongContext,
260    SlidingWindow,
261    Quantization,
262    LoRA,
263}
264
265/// Memory requirements for a model
266#[derive(Debug, Clone)]
267pub struct MemoryRequirements {
268    pub gpu_memory_bytes: usize,
269    pub cpu_memory_bytes: usize,
270    pub kv_cache_bytes: usize,
271    pub peak_memory_bytes: usize,
272    pub fragmentation_overhead: f32,
273}
274
275/// Model inputs for forward pass
276#[derive(Debug, Clone)]
277pub enum ModelInputs {
278    Text {
279        input_ids: Tensor,
280        attention_mask: Option<Tensor>,
281        position_ids: Option<Tensor>,
282    },
283    Image {
284        image_tensor: Tensor,
285        image_features: Option<Tensor>,
286    },
287    Multimodal {
288        input_ids: Tensor,
289        attention_mask: Option<Tensor>,
290        image_tensor: Option<Tensor>,
291        token_type_ids: Option<Tensor>,
292    },
293    Audio {
294        input_features: Tensor,
295        decoder_input_ids: Option<Tensor>,
296    },
297}
298
299impl ModelInputs {
300    pub fn text(text: &str) -> Self {
301        // Placeholder - needs tokenization
302        let input_ids = Tensor::zeros(&[1, text.len()]);
303        Self::Text {
304            input_ids,
305            attention_mask: None,
306            position_ids: None,
307        }
308    }
309
310    pub fn image(image_tensor: Tensor) -> Self {
311        Self::Image {
312            image_tensor,
313            image_features: None,
314        }
315    }
316}
317
318/// Model outputs from forward pass
319#[derive(Debug, Clone)]
320pub enum ModelOutputs {
321    Logits(Tensor),
322    Embeddings(Tensor),
323    ClassificationLogits(Tensor),
324    SequenceClassifierOutput {
325        logits: Tensor,
326        hidden_states: Option<Vec<Tensor>>,
327    },
328    CausalLMOutput {
329        logits: Tensor,
330        past_key_values: Option<Vec<Tensor>>,
331        hidden_states: Option<Vec<Tensor>>,
332        attentions: Option<Vec<Tensor>>,
333    },
334    Seq2SeqLMOutput {
335        logits: Tensor,
336        encoder_last_hidden_state: Tensor,
337        past_key_values: Option<Vec<Tensor>>,
338    },
339}
340
341impl ModelOutputs {
342    pub fn text(&self) -> Option<String> {
343        // Convert logits to text - placeholder implementation
344        match self {
345            ModelOutputs::Logits(_) => Some("Generated text".to_string()),
346            ModelOutputs::CausalLMOutput { .. } => Some("Causal LM output".to_string()),
347            _ => None,
348        }
349    }
350
351    pub fn logits(&self) -> Option<&Tensor> {
352        match self {
353            ModelOutputs::Logits(tensor) => Some(tensor),
354            ModelOutputs::ClassificationLogits(tensor) => Some(tensor),
355            ModelOutputs::CausalLMOutput { logits, .. } => Some(logits),
356            ModelOutputs::Seq2SeqLMOutput { logits, .. } => Some(logits),
357            ModelOutputs::SequenceClassifierOutput { logits, .. } => Some(logits),
358            _ => None,
359        }
360    }
361}
362
363/// Container for loaded model weights
364#[derive(Debug, Clone)]
365pub struct ModelWeights {
366    /// Tensor data indexed by layer name
367    pub tensors: HashMap<String, Tensor>,
368    /// Metadata about the weights
369    pub metadata: WeightMetadata,
370}
371
372/// Metadata about loaded weights
373#[derive(Debug, Clone)]
374pub struct WeightMetadata {
375    /// Total size in bytes
376    pub total_size: usize,
377    /// Number of parameters
378    pub num_parameters: u64,
379    /// Weight format
380    pub format: String,
381    /// Precision info
382    pub dtype: DataType,
383}
384
385/// Model output structure
386#[derive(Debug, Clone)]
387pub struct ModelOutput {
388    pub logits: Tensor,
389    pub hidden_states: Option<Vec<Tensor>>,
390    pub attention_weights: Option<Vec<Tensor>>,
391    pub kv_cache_states: Option<HashMap<String, Tensor>>,
392    pub auxiliary_outputs: HashMap<String, Tensor>,
393}
394
395/// Prepared inputs for model execution
396#[derive(Debug, Clone)]
397pub struct PreparedInputs {
398    pub input_ids: Tensor,
399    pub attention_mask: Option<Tensor>,
400    pub position_ids: Option<Tensor>,
401    pub input_embeddings: Option<Tensor>,
402    pub auxiliary_inputs: HashMap<String, Tensor>,
403}
404
405/// Error types for model operations
406#[derive(Debug)]
407pub enum ModelError {
408    InitializationFailed(String),
409    ComputationFailed(String),
410    InvalidInput(String),
411    DeviceError(String),
412    MemoryError(String),
413    UnsupportedOperation(String),
414    GenerationFailed(String),
415    ValidationFailed(String),
416    ServerError(String),
417    LoadingError(String),
418    ConfigurationError(String),
419    NetworkError(String),
420    FileSystemError(String),
421    CommunicationError(String),
422    RuntimeError(String),
423}
424
425impl fmt::Display for ModelError {
426    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
427        match self {
428            ModelError::InitializationFailed(msg) => write!(f, "Initialization failed: {}", msg),
429            ModelError::ComputationFailed(msg) => write!(f, "Computation failed: {}", msg),
430            ModelError::InvalidInput(msg) => write!(f, "Invalid input: {}", msg),
431            ModelError::DeviceError(msg) => write!(f, "Device error: {}", msg),
432            ModelError::MemoryError(msg) => write!(f, "Memory error: {}", msg),
433            ModelError::UnsupportedOperation(msg) => write!(f, "Unsupported operation: {}", msg),
434            ModelError::GenerationFailed(msg) => write!(f, "Generation failed: {}", msg),
435            ModelError::ValidationFailed(msg) => write!(f, "Validation failed: {}", msg),
436            ModelError::ServerError(msg) => write!(f, "Server error: {}", msg),
437            ModelError::LoadingError(msg) => write!(f, "Loading error: {}", msg),
438            ModelError::ConfigurationError(msg) => write!(f, "Configuration error: {}", msg),
439            ModelError::NetworkError(msg) => write!(f, "Network error: {}", msg),
440            ModelError::FileSystemError(msg) => write!(f, "File system error: {}", msg),
441            ModelError::CommunicationError(msg) => write!(f, "Communication error: {}", msg),
442            ModelError::RuntimeError(msg) => write!(f, "Runtime error: {}", msg),
443        }
444    }
445}
446
447impl std::error::Error for ModelError {}
448
449/// Result type for model operations
450pub type ModelResult<T> = Result<T, ModelError>;
451
452/// Model format types
453#[derive(Debug, Clone, Copy, PartialEq, Eq)]
454pub enum ModelFormat {
455    SafeTensors,
456    GGUF,
457    PyTorch,
458    HuggingFace,
459}
460
461/// Model precision types
462#[derive(Debug, Clone, Copy, PartialEq, Eq)]
463pub enum ModelPrecision {
464    Float32,
465    Float16,
466    BFloat16,
467    Int8,
468    Int4,
469}
470
471impl fmt::Display for ModelPrecision {
472    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
473        match self {
474            ModelPrecision::Float32 => write!(f, "fp32"),
475            ModelPrecision::Float16 => write!(f, "fp16"),
476            ModelPrecision::BFloat16 => write!(f, "bf16"),
477            ModelPrecision::Int8 => write!(f, "int8"),
478            ModelPrecision::Int4 => write!(f, "int4"),
479        }
480    }
481}
482
483/// Model weight metadata
484#[derive(Debug, Clone)]
485pub struct ModelWeightMetadata {
486    pub num_parameters: u64,
487    pub precision: ModelPrecision,
488    pub total_size_bytes: u64,
489    pub shard_count: usize,
490    pub architecture: String,
491    pub model_type: String,
492}
493
494
495/// Normalization types
496#[derive(Debug, Clone, Copy, PartialEq, Eq)]
497pub enum NormalizationType {
498    LayerNorm,
499    RMSNorm,
500    GroupNorm,
501}
502
503/// Position embedding types
504#[derive(Debug, Clone, Copy, PartialEq, Eq)]
505pub enum PositionEmbeddingType {
506    Absolute,
507    Relative,
508    Rotary,
509    ALiBi,
510}
511
512/// Generation statistics
513#[derive(Debug, Clone)]
514pub struct GenerationStats {
515    pub prompt_tokens: usize,
516    pub completion_tokens: usize,
517    pub total_tokens: usize,
518    pub time_to_first_token_ms: f64,
519    pub tokens_per_second: f64,
520    pub total_time_ms: f64,
521    pub cache_hit_rate: f64,
522    pub memory_usage_mb: f64,
523}
524
525/// Inference inputs
526#[derive(Debug, Clone)]
527pub struct InferenceInputs {
528    pub batch_size: usize,
529    pub sequence_length: usize,
530    pub attention_mask: Option<Vec<bool>>,
531    pub position_ids: Option<Vec<u32>>,
532}
533
534/// Inference output
535#[derive(Debug, Clone)]
536pub struct InferenceOutput {
537    pub text: String,
538    pub logits: Option<Vec<f32>>,
539    pub hidden_states: Option<Vec<Tensor>>,
540    pub attention_weights: Option<Vec<Tensor>>,
541    pub generation_stats: Option<GenerationStats>,
542}
543
544/// Calculate strides for a given shape (row-major order)
545pub fn calculate_strides(shape: &[usize]) -> Vec<usize> {
546    let mut strides = vec![1; shape.len()];
547    for i in (0..shape.len() - 1).rev() {
548        strides[i] = strides[i + 1] * shape[i + 1];
549    }
550    strides
551}
552
553/// Supported weight file formats for model loading
554#[derive(Debug, Clone, PartialEq, Eq, Hash)]
555pub enum WeightFormat {
556    /// SafeTensors format (recommended)
557    SafeTensors,
558    /// PyTorch pickle format
559    PyTorch,
560    /// GGUF quantized format
561    GGUF,
562    /// Sharded format (contains base format)
563    Sharded(Box<WeightFormat>),
564}
565
566impl WeightFormat {
567    pub fn from_path(path: &std::path::Path) -> Option<Self> {
568        let extension = path.extension()?.to_str()?;
569        match extension {
570            "safetensors" => Some(Self::SafeTensors),
571            "bin" | "pt" => Some(Self::PyTorch),
572            "gguf" => Some(Self::GGUF),
573            _ => None,
574        }
575    }
576
577    pub fn as_ref(&self) -> &WeightFormat {
578        match self {
579            Self::Sharded(base) => base.as_ref(),
580            other => other,
581        }
582    }
583}
584
585#[cfg(test)]
586mod tests {
587    use super::*;
588
589    #[test]
590    fn test_data_type_size_bytes() {
591        assert_eq!(DataType::Float32.size_bytes(), 4);
592        assert_eq!(DataType::Float16.size_bytes(), 2);
593        assert_eq!(DataType::BFloat16.size_bytes(), 2);
594        assert_eq!(DataType::Int32.size_bytes(), 4);
595        assert_eq!(DataType::Int64.size_bytes(), 8);
596        assert_eq!(DataType::Int8.size_bytes(), 1);
597        assert_eq!(DataType::Int4.size_bytes(), 1);
598        assert_eq!(DataType::Bool.size_bytes(), 1);
599    }
600
601    #[test]
602    fn test_device_display() {
603        assert_eq!(Device::CPU.to_string(), "cpu");
604        assert_eq!(Device::CUDA(0).to_string(), "cuda:0");
605        assert_eq!(Device::CUDA(5).to_string(), "cuda:5");
606        assert_eq!(Device::ROCm(2).to_string(), "rocm:2");
607        assert_eq!(Device::Intel(1).to_string(), "intel:1");
608        assert_eq!(Device::Metal(0).to_string(), "metal:0");
609    }
610
611    #[test]
612    fn test_tensor_creation() {
613        let tensor = Tensor::new(vec![2, 3, 4], DataType::Float32, Device::CPU);
614        assert_eq!(tensor.shape, vec![2, 3, 4]);
615        assert_eq!(tensor.dtype, DataType::Float32);
616        assert_eq!(tensor.device, Device::CPU);
617        assert_eq!(tensor.numel(), 24);
618        assert_eq!(tensor.size_bytes(), 96); // 24 * 4 bytes
619        assert_eq!(tensor.strides, vec![12, 4, 1]); // Row-major strides
620    }
621
622    #[test]
623    fn test_tensor_numel() {
624        let tensor1 = Tensor::new(vec![5], DataType::Float32, Device::CPU);
625        assert_eq!(tensor1.numel(), 5);
626
627        let tensor2 = Tensor::new(vec![2, 3], DataType::Float32, Device::CPU);
628        assert_eq!(tensor2.numel(), 6);
629
630        let tensor3 = Tensor::new(vec![2, 3, 4], DataType::Float32, Device::CPU);
631        assert_eq!(tensor3.numel(), 24);
632    }
633
634    #[test]
635    fn test_tensor_size_bytes() {
636        let tensor_f32 = Tensor::new(vec![10], DataType::Float32, Device::CPU);
637        assert_eq!(tensor_f32.size_bytes(), 40);
638
639        let tensor_f16 = Tensor::new(vec![10], DataType::Float16, Device::CPU);
640        assert_eq!(tensor_f16.size_bytes(), 20);
641
642        let tensor_i8 = Tensor::new(vec![10], DataType::Int8, Device::CPU);
643        assert_eq!(tensor_i8.size_bytes(), 10);
644    }
645
646    #[test]
647    fn test_model_config_creation() {
648        let config = ModelConfig {
649            model_name: "test_model".to_string(),
650            model_path: "/path/to/model".to_string(),
651            max_sequence_length: 2048,
652            vocabulary_size: 32000,
653            num_layers: 32,
654            num_heads: 32,
655            head_dim: 128,
656            hidden_size: 4096,
657            intermediate_size: 11008,
658            dtype: DataType::Float16,
659        };
660
661        assert_eq!(config.model_name, "test_model");
662        assert_eq!(config.max_sequence_length, 2048);
663        assert_eq!(config.vocabulary_size, 32000);
664        assert_eq!(config.dtype, DataType::Float16);
665    }
666
667    #[test]
668    fn test_model_error_display() {
669        let error = ModelError::InitializationFailed("test error".to_string());
670        assert_eq!(error.to_string(), "Initialization failed: test error");
671
672        let error = ModelError::ComputationFailed("math error".to_string());
673        assert_eq!(error.to_string(), "Computation failed: math error");
674
675        let error = ModelError::InvalidInput("bad input".to_string());
676        assert_eq!(error.to_string(), "Invalid input: bad input");
677    }
678
679    #[test]
680    fn test_memory_requirements() {
681        let req = MemoryRequirements {
682            gpu_memory_bytes: 1000000,
683            cpu_memory_bytes: 500000,
684            kv_cache_bytes: 200000,
685            peak_memory_bytes: 1500000,
686            fragmentation_overhead: 0.2,
687        };
688
689        assert_eq!(req.gpu_memory_bytes, 1000000);
690        assert_eq!(req.fragmentation_overhead, 0.2);
691    }
692}
runtime/types.rs

runtime/
types.rs
