Function dequantize_fp8_blockwise

cake-core/src/utils/fp8.rs:43–62

Dequantize a 2-D FP8 weight tensor using its per-block scale factor.

(weight: &Tensor, scale_inv: &Tensor)

/// Dequantize a 2-D FP8 weight tensor using its per-block scale factor.
pub fn dequantize_fp8_blockwise(weight: &Tensor, scale_inv: &Tensor) -> candle_core::Result<Tensor> {
    let (m, n) = weight.dims2()?;
    let bm = FP8_BLOCK_SIZE;
    let bn = FP8_BLOCK_SIZE;
    let blocks_m = m.div_ceil(bm);
    let blocks_n = n.div_ceil(bn);

    // Cast FP8 → F32 on CPU (candle supports this on CPU)
    let weight_f32 = weight.to_dtype(DType::F32)?;
    let scale_f32 = scale_inv.to_dtype(DType::F32)?;

    // Reshape for block-wise broadcast multiply:
    //   weight: [M, N]         → [blocks_m, bm, blocks_n, bn]
    //   scale:  [blocks_m, blocks_n] → [blocks_m, 1, blocks_n, 1]
    let weight_blocked = weight_f32.reshape((blocks_m, bm, blocks_n, bn))?;
    let scale_blocked = scale_f32.reshape((blocks_m, 1usize, blocks_n, 1usize))?;

    let dequantized = weight_blocked.broadcast_mul(&scale_blocked)?;
    dequantized.reshape((m, n))
}

/// Custom VarBuilder backend that wraps MmapedSafetensors and transparently
/// dequantizes FP8-quantized weight tensors on CPU before moving to the target device.