feat(agent): MemoryConfig::float16 stores half-precision embeddings
The setting was persisted in /meta and otherwise ignored: embeddings
were always written as f32. It now does what it says.
clawhdf5-format:
- `DatasetBuilder::with_f16_data` writes IEEE binary16 (numpy float16),
rounding to nearest-even, and `make_f16_type`.
- `clawhdf5_format::float16` holds the f32 <-> f16 conversions, the one
implementation the writer, the reader and the agent all use. Checked
against the `half` crate on 16.7M f32 values and round-trips all 65536
half values; the h5py interop tests confirm the rounding matches
numpy's bit for bit (4020 values incl. ties, subnormals, overflow).
- Reading little-endian float16 as f32 has a fast path.
clawhdf5-agent:
- A float16 store writes /memory/embeddings as half precision, and
`MemoryCache::half_precision` rounds each embedding as it enters the
cache (save, update, WAL replay, and on load of a store still f32 on
disk), so memory and file agree bit for bit and a store searches the
same before and after a reopen (tested).
- Values beyond +-65504 are refused with the new
`MemoryError::InvalidEntry` rather than stored as infinity, on every
save path; batches are all or nothing, and a rejected ephemeral entry
stays in the ephemeral tier. Breaking for exhaustive matches.
- CLI: `create --float16`. Off by default.
Measured on tank, 384-dim, six runs alternating order, medians
(search_harness --float16-study --full): at 100K the file goes from
154.0 to 80.8 MiB (-48%), checkpoint 752 -> 512 ms, open 300 -> 252 ms;
vector recall@10 against an exact scan and hybrid_search latency do not
change. At 10K open is 3 ms slower. Also a test that h5py opens a whole
agent store, f32 and float16, and decodes every dataset.
Docs: README, BENCHMARKS.md ("float16 embedding storage"), CHANGELOG
(including the h5py interop fixes in the previous commit), CLAUDE.md.
Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
@@ -24,6 +24,7 @@ libaec-sys = { path = "../libaec-sys", version = "0.1", optional = true }
|
||||
pco = { version = "1.0", optional = true }
|
||||
|
||||
[dev-dependencies]
|
||||
half = { workspace = true }
|
||||
serde_json = "1"
|
||||
criterion = { workspace = true }
|
||||
clawhdf5-derive = { path = "../clawhdf5-derive", version = "2.7.0" }
|
||||
|
||||
@@ -1076,6 +1076,21 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
|
||||
) {
|
||||
return Ok(native_le_to_vec::<f32>(raw, count));
|
||||
}
|
||||
// Little-endian half precision (numpy float16): widen directly.
|
||||
if matches!(
|
||||
datatype,
|
||||
Datatype::FloatingPoint {
|
||||
size: 2,
|
||||
byte_order: DatatypeByteOrder::LittleEndian,
|
||||
..
|
||||
}
|
||||
) {
|
||||
let (halves, _) = raw[..count * 2].as_chunks::<2>();
|
||||
return Ok(halves
|
||||
.iter()
|
||||
.map(|&b| f16_bits_to_f32(u16::from_le_bytes(b)))
|
||||
.collect());
|
||||
}
|
||||
|
||||
let order = get_byte_order(datatype);
|
||||
let mut result = Vec::with_capacity(count);
|
||||
@@ -1622,36 +1637,7 @@ fn read_f16_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
||||
f16_bits_to_f32(u16::from_le_bytes(buf))
|
||||
}
|
||||
|
||||
/// Convert the bit pattern of an IEEE-754 half (binary16) to an `f32`.
|
||||
fn f16_bits_to_f32(h: u16) -> f32 {
|
||||
let h = h as u32;
|
||||
let sign = (h & 0x8000) << 16;
|
||||
let exp = (h >> 10) & 0x1f;
|
||||
let mant = h & 0x3ff;
|
||||
let bits = if exp == 0 {
|
||||
if mant == 0 {
|
||||
sign // signed zero
|
||||
} else {
|
||||
// Subnormal: normalize into an f32 normal.
|
||||
let mut e: i32 = -1;
|
||||
let mut m = mant;
|
||||
loop {
|
||||
e += 1;
|
||||
m <<= 1;
|
||||
if m & 0x400 != 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let m = m & 0x3ff;
|
||||
sign | (((127 - 15 - e) as u32) << 23) | (m << 13)
|
||||
}
|
||||
} else if exp == 0x1f {
|
||||
sign | 0x7f80_0000 | (mant << 13) // inf / NaN
|
||||
} else {
|
||||
sign | ((exp + (127 - 15)) << 23) | (mant << 13)
|
||||
};
|
||||
f32::from_bits(bits)
|
||||
}
|
||||
use crate::float16::f16_bits_to_f32;
|
||||
|
||||
fn read_f32_bytes(bytes: &[u8], order: &DatatypeByteOrder) -> f32 {
|
||||
let mut buf = [0u8; 4];
|
||||
|
||||
@@ -829,8 +829,12 @@ mod tests {
|
||||
// The HDF5 library rejects a float whose sign position is not inside
|
||||
// its precision; this was hard-coded to 63, so every f32 we wrote was
|
||||
// unreadable by h5py. Byte 2 of the message is the sign position.
|
||||
use crate::type_builders::{make_f32_type, make_f64_type};
|
||||
for (dt, sign) in [(make_f32_type(), 31), (make_f64_type(), 63)] {
|
||||
use crate::type_builders::{make_f16_type, make_f32_type, make_f64_type};
|
||||
for (dt, sign) in [
|
||||
(make_f16_type(), 15),
|
||||
(make_f32_type(), 31),
|
||||
(make_f64_type(), 63),
|
||||
] {
|
||||
let bytes = dt.serialize();
|
||||
assert_eq!(bytes[2], sign, "{dt:?}");
|
||||
let (parsed, _) = Datatype::parse(&bytes).unwrap();
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
//! IEEE-754 half precision (binary16) conversions.
|
||||
//!
|
||||
//! Pure integer bit manipulation, so it works under `no_std` and needs no
|
||||
//! `libm`. The writer ([`crate::type_builders::DatasetBuilder::with_f16_data`]),
|
||||
//! the reader and `clawhdf5-agent`'s half-precision embedding store all use
|
||||
//! these two functions, so a value rounded in memory is bit-for-bit the value
|
||||
//! that reads back from the file.
|
||||
|
||||
/// Largest finite half-precision value. Anything larger in magnitude rounds
|
||||
/// to infinity.
|
||||
pub const F16_MAX: f32 = 65504.0;
|
||||
|
||||
/// Convert an `f32` to the bit pattern of the nearest half-precision value,
|
||||
/// rounding ties to even (the IEEE default, and what numpy and the `half`
|
||||
/// crate do).
|
||||
///
|
||||
/// Values beyond ±[`F16_MAX`] become ±infinity, values too small for a
|
||||
/// subnormal become signed zero, and NaN stays NaN (quiet, payload
|
||||
/// truncated).
|
||||
pub fn f32_to_f16_bits(value: f32) -> u16 {
|
||||
let x = value.to_bits();
|
||||
let sign = (x >> 16) & 0x8000;
|
||||
let exp = x & 0x7F80_0000;
|
||||
let man = x & 0x007F_FFFF;
|
||||
|
||||
// Infinity and NaN.
|
||||
if exp == 0x7F80_0000 {
|
||||
let quiet_nan = if man == 0 { 0 } else { 0x0200 };
|
||||
return (sign | 0x7C00 | quiet_nan | (man >> 13)) as u16;
|
||||
}
|
||||
|
||||
let half_exp = ((exp >> 23) as i32) - 127 + 15;
|
||||
|
||||
// Too large: infinity.
|
||||
if half_exp >= 0x1F {
|
||||
return (sign | 0x7C00) as u16;
|
||||
}
|
||||
|
||||
// Subnormal half, or zero.
|
||||
if half_exp <= 0 {
|
||||
if 14 - half_exp > 24 {
|
||||
return sign as u16;
|
||||
}
|
||||
let man = man | 0x0080_0000; // implicit leading bit
|
||||
let shift = (14 - half_exp) as u32;
|
||||
let mut half_man = man >> shift;
|
||||
let round_bit = 1u32 << (shift - 1);
|
||||
// Round half to even: up if above half, or exactly half and odd.
|
||||
if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 {
|
||||
half_man += 1;
|
||||
}
|
||||
return (sign | half_man) as u16;
|
||||
}
|
||||
|
||||
// Normal half. A mantissa carry correctly rolls into the exponent (and
|
||||
// from the largest finite value into infinity).
|
||||
let half = sign | ((half_exp as u32) << 10) | (man >> 13);
|
||||
let round_bit = 0x0000_1000;
|
||||
if (man & round_bit) != 0 && (man & (3 * round_bit - 1)) != 0 {
|
||||
(half + 1) as u16
|
||||
} else {
|
||||
half as u16
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert the bit pattern of a half-precision value to `f32` (exact: every
|
||||
/// half value is representable as an `f32`).
|
||||
pub fn f16_bits_to_f32(h: u16) -> f32 {
|
||||
let h = h as u32;
|
||||
let sign = (h & 0x8000) << 16;
|
||||
let exp = (h >> 10) & 0x1f;
|
||||
let mant = h & 0x3ff;
|
||||
let bits = if exp == 0 {
|
||||
if mant == 0 {
|
||||
sign // signed zero
|
||||
} else {
|
||||
// Subnormal: normalize into an f32 normal.
|
||||
let mut e: i32 = -1;
|
||||
let mut m = mant;
|
||||
loop {
|
||||
e += 1;
|
||||
m <<= 1;
|
||||
if m & 0x400 != 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let m = m & 0x3ff;
|
||||
sign | (((127 - 15 - e) as u32) << 23) | (m << 13)
|
||||
}
|
||||
} else if exp == 0x1f {
|
||||
sign | 0x7f80_0000 | (mant << 13) // inf / NaN
|
||||
} else {
|
||||
sign | ((exp + 127 - 15) << 23) | (mant << 13)
|
||||
};
|
||||
f32::from_bits(bits)
|
||||
}
|
||||
|
||||
/// Round an `f32` to the nearest half-precision value, returned as `f32`.
|
||||
pub fn round_to_f16(value: f32) -> f32 {
|
||||
f16_bits_to_f32(f32_to_f16_bits(value))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn every_half_value_round_trips() {
|
||||
for bits in 0..=u16::MAX {
|
||||
let v = f16_bits_to_f32(bits);
|
||||
if v.is_nan() {
|
||||
assert!(f16_bits_to_f32(f32_to_f16_bits(v)).is_nan(), "{bits:#06x}");
|
||||
} else {
|
||||
assert_eq!(f32_to_f16_bits(v), bits, "{bits:#06x} -> {v}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_the_half_crate() {
|
||||
// Every 257th f32 bit pattern (~16.7M values) covers every exponent,
|
||||
// the subnormal range, both signs, ties and the overflow boundary.
|
||||
let mut bits: u32 = 0;
|
||||
loop {
|
||||
let v = f32::from_bits(bits);
|
||||
let ours = f32_to_f16_bits(v);
|
||||
let theirs = half::f16::from_f32(v);
|
||||
if v.is_nan() {
|
||||
assert!(theirs.is_nan() && f16_bits_to_f32(ours).is_nan());
|
||||
} else {
|
||||
assert_eq!(ours, theirs.to_bits(), "{bits:#010x} ({v:e})");
|
||||
assert_eq!(f16_bits_to_f32(ours).to_bits(), theirs.to_f32().to_bits());
|
||||
}
|
||||
match bits.checked_add(257) {
|
||||
Some(b) => bits = b,
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rounds_ties_to_even_and_saturates_to_infinity() {
|
||||
// 1 + 2^-11 is exactly halfway between 1.0 and the next half (1 + 2^-10).
|
||||
assert_eq!(round_to_f16(1.0 + 2f32.powi(-11)), 1.0);
|
||||
assert_eq!(
|
||||
round_to_f16(1.0 + 3.0 * 2f32.powi(-11)),
|
||||
1.0 + 2.0 * 2f32.powi(-10)
|
||||
);
|
||||
assert_eq!(round_to_f16(F16_MAX), F16_MAX);
|
||||
assert_eq!(round_to_f16(65520.0), f32::INFINITY); // halfway to 2^16 rounds up
|
||||
assert_eq!(round_to_f16(-1e9), f32::NEG_INFINITY);
|
||||
assert_eq!(round_to_f16(1e-9).to_bits(), 0);
|
||||
assert_eq!(round_to_f16(-1e-9).to_bits(), (-0.0f32).to_bits());
|
||||
}
|
||||
}
|
||||
@@ -72,6 +72,7 @@ pub mod filter_pipeline;
|
||||
pub mod filters;
|
||||
mod filters_szip;
|
||||
pub mod fixed_array;
|
||||
pub mod float16;
|
||||
pub mod fractal_heap;
|
||||
pub mod global_heap;
|
||||
pub mod group_info;
|
||||
|
||||
@@ -56,6 +56,21 @@ pub fn make_f64_type() -> Datatype {
|
||||
}
|
||||
}
|
||||
|
||||
/// IEEE-754 half precision (binary16), little-endian — numpy's `float16`.
|
||||
pub fn make_f16_type() -> Datatype {
|
||||
Datatype::FloatingPoint {
|
||||
size: 2,
|
||||
byte_order: DatatypeByteOrder::LittleEndian,
|
||||
bit_offset: 0,
|
||||
bit_precision: 16,
|
||||
exponent_location: 10,
|
||||
exponent_size: 5,
|
||||
mantissa_location: 0,
|
||||
mantissa_size: 10,
|
||||
exponent_bias: 15,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn make_f32_type() -> Datatype {
|
||||
Datatype::FloatingPoint {
|
||||
size: 4,
|
||||
@@ -478,6 +493,24 @@ impl DatasetBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Store `data` as IEEE half precision (numpy `float16`), rounding each
|
||||
/// value to the nearest half ([`crate::float16::f32_to_f16_bits`]).
|
||||
/// Half the bytes of [`Self::with_f32_data`], at about three significant
|
||||
/// decimal digits; values beyond ±65504 become ±infinity. Reading it back
|
||||
/// with `read_f32` yields the rounded values exactly.
|
||||
pub fn with_f16_data(&mut self, data: &[f32]) -> &mut Self {
|
||||
self.datatype = Some(make_f16_type());
|
||||
let mut b = Vec::with_capacity(data.len() * 2);
|
||||
for &v in data {
|
||||
b.extend_from_slice(&crate::float16::f32_to_f16_bits(v).to_le_bytes());
|
||||
}
|
||||
self.data = Some(b);
|
||||
if self.shape.is_none() {
|
||||
self.shape = Some(vec![data.len() as u64]);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_i32_data(&mut self, data: &[i32]) -> &mut Self {
|
||||
self.datatype = Some(make_i32_type());
|
||||
let mut b = Vec::with_capacity(data.len() * 4);
|
||||
|
||||
Reference in New Issue
Block a user