MCPcopy Create free account
hub / github.com/evilsocket/cake / load_tensor

Method load_tensor

cake-core/src/utils/fp8.rs:101–130  ·  view source on GitHub ↗
(
        &self,
        name: &str,
        dtype: DType,
        dev: &Device,
    )

Source from the content-addressed store, hash-verified

99
100impl Fp8Backend {
101 fn load_tensor(
102 &self,
103 name: &str,
104 dtype: DType,
105 dev: &Device,
106 ) -> candle_core::Result<Tensor> {
107 let scale_name = format!("{name}_scale_inv");
108
109 if self.inner.get(&scale_name).is_ok() {
110 // FP8 quantized tensor — dequantize on CPU then move to device
111 let weight = self.inner.load(name, &Device::Cpu)?;
112 let scale = self.inner.load(&scale_name, &Device::Cpu)?;
113
114 let dequantized = dequantize_fp8_blockwise(&weight, &scale)?;
115 dequantized.to_dtype(dtype)?.to_device(dev)
116 } else {
117 // Non-quantized tensor — check if the on-file dtype needs CPU-side handling
118 let view = self.inner.get(name)?;
119 let file_dtype: DType = view.dtype().try_into()?;
120
121 if file_dtype == DType::F8E4M3 {
122 // FP8 without scale (shouldn't happen, but handle gracefully)
123 let tensor = self.inner.load(name, &Device::Cpu)?;
124 tensor.to_dtype(dtype)?.to_device(dev)
125 } else {
126 // Normal path — load directly on target device
127 self.inner.load(name, dev)?.to_dtype(dtype)
128 }
129 }
130 }
131}
132
133/// Create a VarBuilder that transparently dequantizes FP8 weights.

Callers 2

getMethod · 0.45
get_uncheckedMethod · 0.45

Calls 4

dtypeMethod · 0.80
dequantize_fp8_blockwiseFunction · 0.70
getMethod · 0.45
loadMethod · 0.45

Tested by

no test coverage detected