Java · Linear Algebra

Linear Algebra

3.0.0-rc.1 · 4c3e5023

§ Sparse Multiplication

multiply is the linear-algebra sibling of search: it reports every pair the traversal scored rather than only the nearest, as a ragged product with one run of pairs per left row, in traversal order. It runs on the inner-product kinds alone, and a non-zero wanted count keeps the first pairs each row’s walk scored. The keys, products and per-row offsets all land in the arena, so a steady loop allocates none.

DenseProducts products = index.on(pool).multiply(left, arena);
for (int row = 0; row < products.queries(); ++row) {
    for (long rank = 0; rank < products.count(row); ++rank) {
        System.out.println(products.key(row, rank) + " x " + products.score(row, rank));
}
}

§ Late Interaction

maxsim scores multi-row queries against multi-vector documents: 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, so valuesPerKey bounds how many rows share a key and the wanted count is the “k” over documents rather than over rows. Every query carries the same number of rows, so a padded token batch applies as is. Each query’s keys land largest score first in the same ragged DenseProducts that multiply returns, and a wanted count of zero keeps every key the query scored.

DenseOptions tokens = new DenseOptions().dimensions(128).metric(DenseMetric.INNER).valuesPerKey(512);
try (DenseIndex index = DenseIndex.initRam(tokens, new int[0])) {
    DenseProducts documents = index.on(pool).maxsim(queryTokens, 32, 10, arena);
    for (long rank = 0; rank < documents.count(0); ++rank) {
        System.out.println(documents.key(0, rank) + " scores " + documents.score(0, rank));
    }
}

It runs on the inner-product kinds alone, since a maximum over logits is what late interaction sums; a graph frozen under a distance kind reports Status.CONFIG_MISMATCH.

§ 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 the arena already divided by the partition. The index stores no values, so a ValueLookup names the row of a values matrix each scored key resolves to. It borrows that matrix and builds its upcall once, so every call reuses both: a native matrix is read in place, a heap one is copied at construction. The library states the scalar and width it reads each row as, and a values matrix of another kind, a row outside it, or a lookup that throws leaves the key without a row, which refuses the call with Status.SHAPE_MISMATCH; the first exception the lookup raised outranks that status, and is raised once the call is final. The lookup runs on the executor’s threads together, so it must tolerate concurrent calls. It is a host function, so it serves RAM handles alone: a VRAM walk resolves rows through a device function, which this binding does not build, and an attend into a VRAM arena refuses with Status.DEVICE.

ValueLookup cache = new ValueLookup(RowsView.of(valueRows, 128, Scalar.BF16), rowOf::get);
DenseSoftmaxState[] carry = new DenseSoftmaxState[queries.rows()];
DenseVectors attended = index.on(pool).attend(queries, cache, new DenseAttendConfig().headCount(8).carry(carry), arena);

With headCount(8), row t * 8 + h of the queries is head h of token t, and the heads of one token share its walk. The carry holds one DenseSoftmaxState per query row, lent by the caller rather than placed in the arena, which folds this branch into another over the same queries, such as a recent window kept out of the index. DenseAttendConfig.outputsScalar picks the scalar the outputs land as, the queries’ own by default.