(&self, tensor: &Tensor)
| 591 | } |
| 592 | |
| 593 | fn get_or_upload(&self, tensor: &Tensor) -> Result<Arc<MappedBuffer>> { |
| 594 | let id = tensor.id(); |
| 595 | // Check activation cache first (ephemeral outputs from recent dispatches) |
| 596 | { |
| 597 | let mut act_cache = self.activation_cache.lock().unwrap(); |
| 598 | if let Some(pos) = act_cache.iter().position(|(tid, _)| *tid == id) { |
| 599 | let (_, buf) = act_cache.remove(pos); |
| 600 | return Ok(buf); |
| 601 | } |
| 602 | } |
| 603 | // Check weight cache by TensorId (fast path for stable tensors) |
| 604 | { |
| 605 | let cache = self.weight_cache.lock().unwrap(); |
| 606 | if let Some(buf) = cache.buffers.get(&id) { |
| 607 | return Ok(buf.clone()); |
| 608 | } |
| 609 | } |
| 610 | // View cache: only for NON-CONTIGUOUS tensors (weight.t() views). |
| 611 | // Contiguous tensors (activations) must NOT use the view cache because |
| 612 | // freed tensors can be reallocated at the same address, causing stale hits. |
| 613 | let is_view = !tensor.is_contiguous(); |
| 614 | let vk = if is_view { Self::view_key(tensor) } else { None }; |
| 615 | if let Some(ref key) = vk { |
| 616 | let cache = self.weight_cache.lock().unwrap(); |
| 617 | if let Some(buf) = cache.views.get(key) { |
| 618 | return Ok(buf.clone()); |
| 619 | } |
| 620 | } |
| 621 | // Upload: convert to f32 contiguous, copy to GPU |
| 622 | let tensor = if tensor.dtype() == DType::F32 { tensor.clone() } else { tensor.to_dtype(DType::F32)? }; |
| 623 | let tensor = if tensor.is_contiguous() { tensor } else { tensor.contiguous()? }; |
| 624 | let n = tensor.elem_count(); |
| 625 | let bytes = (n * 4) as u64; |
| 626 | let mut alloc_guard = self.allocator.lock().unwrap(); |
| 627 | let alloc = alloc_guard.as_mut().unwrap(); |
| 628 | let buf = Self::alloc_mapped_buffer( |
| 629 | &self.vk_device, |
| 630 | alloc, |
| 631 | bytes, |
| 632 | vk::BufferUsageFlags::STORAGE_BUFFER, |
| 633 | self.uma_memory_type, |
| 634 | ) |
| 635 | .map_err(candle_core::Error::Msg)?; |
| 636 | // Write directly from tensor storage — avoids intermediate Vec allocation |
| 637 | let (storage, layout) = tensor.storage_and_layout(); |
| 638 | if let candle_core::Storage::Cpu(cpu) = &*storage { |
| 639 | let slice: &[f32] = cpu.as_slice()?; |
| 640 | let offset = layout.start_offset(); |
| 641 | buf.write_f32(&slice[offset..offset + n]); |
| 642 | } else { |
| 643 | drop(storage); |
| 644 | let data = Self::to_f32_vec(&tensor)?; |
| 645 | buf.write_f32(&data); |
| 646 | } |
| 647 | let buf = Arc::new(buf); |
| 648 | let mut cache = self.weight_cache.lock().unwrap(); |
| 649 | if let Some(key) = vk { |
| 650 | // Non-contiguous view (e.g., weight.t()): store in view cache only. |
no test coverage detected