For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo struct
QSABlockScore
struct QSABlockScore
Relu-summed, block-causally masked QSA block scores.
Tensor shapes: - scores : [queries, out_cols] (OUT, f32) - q : [queries, num_heads, head_dim] - block_keys : [batch * max_blocks, head_dim] - query_positions : [queries] (int32) - input_row_offsets : [batch + 1] (uint32)
block_keys holds each sequence's blocks at the uniform stride
max_blocks = rows // batch, which is exactly what qsa_compress_keys
writes -- score column j is sequence-local block j, so a selection
expands to sequence-local token indices for the consumer.
out_cols must be at least max_blocks; the surplus is filled with -inf
so a consumer can block-top-k a fixed-width row even when the live block
count is below its k.
Implemented traits
Methods
execute
static def execute[dtype: DType, target: StringSpan[ImmStaticOrigin], num_heads: Int, head_dim: Int, ratio: Int, score_scale: StringSpan[ImmStaticOrigin]](scores: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=scores.static_spec], q: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=q.static_spec], block_keys: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=block_keys.static_spec], query_positions: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=query_positions.static_spec], input_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec], ctx: DeviceContext)