def attend[queries_origin: ImmOrigin, carry_origin: MutOrigin, stats_origin: MutOrigin, //, dtype: DType](
self,
queries: RowsView[queries_origin],
value_lookup: def(Int, Pointer[Key, ImmUntrackedOrigin], UInt, ScalarKind, UInt, Pointer[Int, MutAnyOrigin]) abi("C") thin -> None,
value_lookup_state: Int,
mut arena: Arena,
config: DenseAttendConfig[carry_origin, stats_origin] = DenseAttendConfig(head_count=Int(0), carry=Span(), expansion_factor=Int(0), score_scale=Float32("0"), stats=Optional())
) -> DenseVectors[dtype, origin_of(arena)]
Every key a walk scores folds into each head's running maximum, partition, and weighted
value sum, as Flash Attention folds key tiles, with no top-k cut: the long tail the walk
scores is what brings the result toward exact dense attention. The outputs come back
already divided by the partition.
Row t · head_count + h of queries, the outputs, and carry is head h of token t. One
index holds one layer's one key-value head, so the head_count query heads of grouped-query
attention share its keys and one walk per token.
The index stores no values: value_lookup resolves the keys of every scored batch to their
value rows, so only keys a walk scores are ever resolved. carry is the softmax's m and
l, one DenseSoftmaxState per query row, lent by the caller: supplying it folds this branch
into another over the same queries, such as the recent window kept out of the index, by
partial-softmax rescaling.
Parameters
Returns
One attended row per query row, as wide as the keys, read in place from the arena.
Raises
Error with
SHAPE_MISMATCH if
dtype names no scalar kind,
head_count does not divide the rows, the carry holds fewer states than there are rows, or a lookup yields a null row; with
CONFIG_MISMATCH if the graph was frozen under a distance kind.