Linear Algebra

§ Sparse Multiplication

multiply is a sparse matrix product: every pair the walk scored for each left row, in traversal order, lands in an arena as one ragged run, with one more offset than left rows, as Apache Arrow lays a ragged batch out.

var products = index.on(executor).multiply(left, arena, DenseMultiplyConfig(wanted=64))
var offsets = products.offsets()
for row in range(left_rows):
    print(row, Int(offsets[row + 1] - offsets[row]), "pairs")

§ 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. The keys land as multiply’s pairs do, query i spanning offsets()[i] up to offsets()[i + 1], largest score first, and a zero wanted keeps every key a query scored.

var documents = index.on(executor).maxsim(tokens, tokens_per_query, 10, arena)  # DenseProducts, as multiply answers
var offsets = documents.offsets()
for query in range(documents.queries):
    print(query, Int(offsets[query + 1] - offsets[query]), "documents")

§ 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 attended rows land in the arena already divided by the partition. The index stores no values, so a ValueLookup resolves every scored batch of keys to value rows, columns values of scalar apiece, the index’s own width and kind, which it checks before answering any row. A key without a row refuses the call with Status.SHAPE_MISMATCH, since a missing row is a wrong softmax rather than a smaller one. A VRAM index calls its lookup inside the device walk, so there value_lookup must be a __device__ function’s address, read with cudaMemcpyFromSymbol on NVIDIA or hipMemcpyFromSymbol on AMD; a host function there is undefined, never refused.

def value_rows(
    state: Address, keys: Pointer[Key, ImmUntrackedOrigin], count: c_size_t,
    scalar: ScalarKind, columns: c_size_t, rows: Pointer[Address, MutAnyOrigin],
) abi("C"):
    ...  # write one row address per key, or NULL_ADDRESS when it has none

var carry = List[DenseSoftmaxState](length=queries_count, fill=DenseSoftmaxState())
var context = index.on(executor).attend[DType.float32](
    queries, value_rows, cache_address, arena, DenseAttendConfig(head_count=8, carry=carry)
)

With head_count=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; a shorter one is refused, since the C structure carries no count for it.