haagenti_neural/
decoder.rs

1//! Neural decoder for decompressing model weights
2
3use crate::{EncodedTensor, LayerCodebook, LayerType, NeuralError, Result};
4use serde::{Deserialize, Serialize};
5
6/// Configuration for the decoder
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct DecoderConfig {
9    /// Whether to apply residual refinement
10    pub apply_residuals: bool,
11    /// Batch size for decoding
12    pub batch_size: usize,
13    /// Use GPU acceleration (future)
14    pub use_gpu: bool,
15    /// Prefetch next layer while decoding
16    pub prefetch: bool,
17}
18
19impl Default for DecoderConfig {
20    fn default() -> Self {
21        Self {
22            apply_residuals: true,
23            batch_size: 1024,
24            use_gpu: false,
25            prefetch: true,
26        }
27    }
28}
29
30/// Decoded tensor result
31#[derive(Debug, Clone)]
32pub struct DecodedTensor {
33    /// Tensor name
34    pub name: String,
35    /// Tensor shape
36    pub shape: Vec<usize>,
37    /// Decoded data
38    pub data: Vec<f32>,
39    /// Decoding statistics
40    pub stats: DecodingStats,
41}
42
43/// Statistics from decoding
44#[derive(Debug, Clone, Default, Serialize, Deserialize)]
45pub struct DecodingStats {
46    /// Time to decode in microseconds
47    pub decode_time_us: u64,
48    /// Codebook lookup time
49    pub lookup_time_us: u64,
50    /// Residual application time
51    pub residual_time_us: u64,
52    /// Number of vectors decoded
53    pub vectors_decoded: usize,
54}
55
56/// Neural decoder for model decompression
57#[derive(Debug)]
58pub struct NeuralDecoder {
59    config: DecoderConfig,
60    codebooks: LayerCodebook,
61}
62
63impl NeuralDecoder {
64    /// Create a new decoder with codebooks
65    pub fn new(config: DecoderConfig, codebooks: LayerCodebook) -> Self {
66        Self { config, codebooks }
67    }
68
69    /// Decode an encoded tensor
70    pub fn decode_tensor(&self, encoded: &EncodedTensor) -> Result<DecodedTensor> {
71        let start = std::time::Instant::now();
72
73        let codebook = self
74            .codebooks
75            .get(encoded.layer_type)
76            .ok_or_else(|| NeuralError::CodebookNotFound(format!("{:?}", encoded.layer_type)))?;
77
78        // Decode indices to vectors
79        let lookup_start = std::time::Instant::now();
80        let mut data = codebook.decode_batch(&encoded.indices);
81        let lookup_time = lookup_start.elapsed();
82
83        // Apply residuals if present
84        let residual_start = std::time::Instant::now();
85        if self.config.apply_residuals {
86            if let Some(residuals) = encoded.residuals.as_ref() {
87                self.apply_residuals(&mut data, residuals, encoded.residual_scale);
88            }
89        }
90        let residual_time = residual_start.elapsed();
91
92        let total_time = start.elapsed();
93
94        let stats = DecodingStats {
95            decode_time_us: total_time.as_micros() as u64,
96            lookup_time_us: lookup_time.as_micros() as u64,
97            residual_time_us: residual_time.as_micros() as u64,
98            vectors_decoded: encoded.indices.len(),
99        };
100
101        Ok(DecodedTensor {
102            name: encoded.name.clone(),
103            shape: encoded.shape.clone(),
104            data,
105            stats,
106        })
107    }
108
109    /// Apply residual refinement
110    fn apply_residuals(&self, data: &mut [f32], residuals: &[i8], scale: f32) {
111        for (d, &r) in data.iter_mut().zip(residuals.iter()) {
112            *d += r as f32 * scale;
113        }
114    }
115
116    /// Decode multiple tensors
117    pub fn decode_tensors(&self, encoded: &[EncodedTensor]) -> Result<Vec<DecodedTensor>> {
118        encoded.iter().map(|e| self.decode_tensor(e)).collect()
119    }
120
121    /// Decode with streaming (yields chunks)
122    pub fn decode_streaming<'a>(
123        &'a self,
124        encoded: &'a EncodedTensor,
125        chunk_size: usize,
126    ) -> impl Iterator<Item = Result<Vec<f32>>> + 'a {
127        let codebook = self.codebooks.get(encoded.layer_type);
128
129        encoded
130            .indices
131            .chunks(chunk_size)
132            .enumerate()
133            .map(move |(i, chunk)| {
134                let cb = codebook.ok_or_else(|| {
135                    NeuralError::CodebookNotFound(format!("{:?}", encoded.layer_type))
136                })?;
137
138                let mut data: Vec<f32> = chunk
139                    .iter()
140                    .flat_map(|&idx| {
141                        cb.get_centroid(idx as usize)
142                            .map(|c| c.to_vec())
143                            .unwrap_or_else(|| vec![0.0; cb.config.centroid_dim])
144                    })
145                    .collect();
146
147                // Apply residuals for this chunk
148                if self.config.apply_residuals {
149                    if let Some(residuals) = encoded.residuals.as_ref() {
150                        let dim = cb.config.centroid_dim;
151                        let start = i * chunk_size * dim;
152                        let end = (start + data.len()).min(residuals.len());
153                        if start < residuals.len() {
154                            self.apply_residuals(
155                                &mut data[..end - start],
156                                &residuals[start..end],
157                                encoded.residual_scale,
158                            );
159                        }
160                    }
161                }
162
163                Ok(data)
164            })
165    }
166
167    /// Get decoder configuration
168    pub fn config(&self) -> &DecoderConfig {
169        &self.config
170    }
171
172    /// Estimate decode throughput in GB/s
173    pub fn estimate_throughput(&self, stats: &DecodingStats) -> f32 {
174        if stats.decode_time_us == 0 {
175            return 0.0;
176        }
177
178        let bytes_decoded = stats.vectors_decoded * 64 * 4; // Assuming 64-dim FP32
179        let seconds = stats.decode_time_us as f64 / 1_000_000.0;
180        (bytes_decoded as f64 / seconds / 1e9) as f32
181    }
182}
183
184/// GPU-accelerated decoder (placeholder for CUDA implementation)
185#[derive(Debug)]
186pub struct GpuDecoder {
187    /// Codebook data on GPU
188    codebook_buffers: Vec<GpuBuffer>,
189    /// Output buffer
190    output_buffer: GpuBuffer,
191}
192
193/// Placeholder for GPU buffer
194#[derive(Debug, Default)]
195struct GpuBuffer {
196    size: usize,
197    device_ptr: u64,
198}
199
200impl GpuDecoder {
201    /// Create GPU decoder (requires CUDA)
202    pub fn new(_codebooks: &LayerCodebook) -> Result<Self> {
203        // In real implementation, this would:
204        // 1. Allocate GPU memory for codebooks
205        // 2. Copy codebook data to GPU
206        // 3. Compile decode kernels
207
208        Ok(Self {
209            codebook_buffers: Vec::new(),
210            output_buffer: GpuBuffer::default(),
211        })
212    }
213
214    /// Decode on GPU
215    pub fn decode_gpu(&self, _indices: &[u16], _layer_type: LayerType) -> Result<Vec<f32>> {
216        // Placeholder - would launch CUDA kernel
217        Err(NeuralError::DecodingError("GPU not available".into()))
218    }
219
220    /// Get total GPU memory allocated (in bytes)
221    pub fn gpu_memory_usage(&self) -> usize {
222        let codebook_mem: usize = self.codebook_buffers.iter().map(|b| b.size).sum();
223        codebook_mem + self.output_buffer.size
224    }
225
226    /// Check if GPU is available
227    pub fn is_available(&self) -> bool {
228        self.output_buffer.device_ptr != 0
229    }
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235    use crate::{EncoderConfig, NeuralEncoder};
236
237    #[test]
238    fn test_decode_tensor() {
239        // First encode something
240        let codebooks = LayerCodebook::with_defaults("test");
241        let encoder = NeuralEncoder::new(EncoderConfig::default(), codebooks.clone());
242
243        let original: Vec<f32> = (0..640).map(|i| (i as f32 / 640.0) - 0.5).collect();
244        let shape = vec![10, 64];
245
246        let encoded = encoder
247            .encode_tensor("test.weight", &original, &shape, LayerType::AttentionQK)
248            .unwrap();
249
250        // Now decode
251        let decoder = NeuralDecoder::new(DecoderConfig::default(), codebooks);
252        let decoded = decoder.decode_tensor(&encoded).unwrap();
253
254        assert_eq!(decoded.name, "test.weight");
255        assert_eq!(decoded.shape, vec![10, 64]);
256        assert_eq!(decoded.data.len(), 640);
257    }
258
259    #[test]
260    fn test_decode_streaming() {
261        let codebooks = LayerCodebook::with_defaults("test");
262        let encoder = NeuralEncoder::new(EncoderConfig::default(), codebooks.clone());
263
264        let original: Vec<f32> = (0..640).map(|i| (i as f32 / 640.0) - 0.5).collect();
265        let encoded = encoder
266            .encode_tensor("test", &original, &[10, 64], LayerType::AttentionQK)
267            .unwrap();
268
269        let decoder = NeuralDecoder::new(DecoderConfig::default(), codebooks);
270
271        let chunks: Vec<_> = decoder
272            .decode_streaming(&encoded, 3)
273            .collect::<Result<Vec<_>>>()
274            .unwrap();
275
276        assert!(!chunks.is_empty());
277    }
278
279    #[test]
280    fn test_residual_application() {
281        let decoder = NeuralDecoder::new(
282            DecoderConfig::default(),
283            LayerCodebook::with_defaults("test"),
284        );
285
286        let mut data = vec![1.0, 2.0, 3.0, 4.0];
287        let residuals = vec![10i8, -10, 20, -20];
288        let scale = 0.01;
289
290        decoder.apply_residuals(&mut data, &residuals, scale);
291
292        assert!((data[0] - 1.1).abs() < 0.001);
293        assert!((data[1] - 1.9).abs() < 0.001);
294        assert!((data[2] - 3.2).abs() < 0.001);
295        assert!((data[3] - 3.8).abs() < 0.001);
296    }
297
298    #[test]
299    fn test_gpu_decoder_construction() {
300        let codebooks = LayerCodebook::with_defaults("test");
301        let decoder = GpuDecoder::new(&codebooks).unwrap();
302
303        // GPU decode should fail gracefully when no GPU is available
304        let result = decoder.decode_gpu(&[0, 1, 2], LayerType::AttentionQK);
305        assert!(result.is_err());
306    }
307}
haagenti_neural/decoder.rs

haagenti_neural/
decoder.rs
