Linear Algebra

§ Sparse Multiplication

multiply, maxsim and attend treat the index as a sparse operand rather than as a kNN store, and each folds its arithmetic into the graph walk instead of materializing a match list. All three need logits rather than distances, so they run on an inner-product index alone.

multiply is a sparse matrix product: every left row is walked against the index, and every pair the walk scored lands in an arena as one ragged run, in traversal order. wanted keeps the first pairs of each row, where zero keeps every pair, and required counts how many there were. The run’s Apache Arrow offsets land in the arena beside the keys and products.

let mut arena = Arena::new_in(&executor)?;
let products = index.on(&executor).multiply(&rows(&left, 768)?, 32, &mut arena)?;
for row in products.iter()? {
    println!("{} pairs", row.keys.len());
}

§ Late Interaction

maxsim is late interaction - every query row takes its best match within a key, and those maxima sum to that key’s score. A document is one key holding its token rows, which is what a values_per_key above one exists for, and the wanted count is over documents rather than rows. Each query lands its best documents, largest score first, in the same ragged DenseProducts run multiply returns, and a wanted of zero keeps every document the query scored.

let documents = index.on(&executor).maxsim(&rows(&tokens, 768)?, tokens_per_query, 10, &mut arena)?;
for query in documents.iter()? {
    println!("best {:?} of {} documents", query.keys.first(), query.keys.len());
}

§ Attention

attend is Flash Attention folded into the walk: every key it scores folds into each head’s running maximum, partition and weighted value sum, with no top-k cut, and the outputs land in an arena already divided by the partition, as rows of the queries’ own scalar type. The index stores no values, so a ValueLookup closure resolves every scored key to its value row, as wide as the keys, stored in the index’s own quantization, and borrowed from storage outliving the call. A None row, or rows of another scalar type or width than the library reads, refuses the call with Status::ShapeMismatch. It runs on the executor’s threads together, hence Sync, and a panic inside it refuses the call. A VRAM walk resolves values on the device, where no Rust closure runs, so an arena from a VRAM executor is refused with Status::Device. With a head_count of 8, row t * 8 + h of the queries is head h of token t, and the heads of one token share its walk. A carry of one DenseSoftmaxState per query row, lent by you, folds this branch into another over the same queries, such as a recent window kept out of the index. score_scale is conventionally the inverse square root of the head width, and zero derives it from the frozen geometry.

let value_lookup: &ValueLookup<bf16> = &|key| cache.get(&key).map(Vec::as_slice);
let mut carry = vec![DenseSoftmaxState::default(); queries_count];
let context = index.on(&executor).attend(&rows(&queries, 128)?, value_lookup, DenseAttendConfig {
    carry: Some(&mut carry),
    ..DenseAttendConfig::from(8)
}, &mut arena)?;
println!("{:?}", context.view()?.shape());