Attention¶
Scaled dot-product attention, grouped-query attention, and paged attention used by the model.
anllm.core.attention.scaled_dot_product_attention_simple(query, key, value, scale=None, mask=None)
¶
A simple implementation of scaled dot product attention. Assuming Q, K, V are of the same shape. Assuming mask is always a float array that you can add to the scores.
anllm.core.attention.scaled_dot_product_attention_grouped(query, key, value, scale=None, mask=None)
¶
Potential input of the mask: - mx.array that can broadcast to B * H_q * L * S, which needs to be reshaped to match multi-head dimensions - None which will be ignored
anllm.core.attention.paged_attention(query, key_pages, value_pages, block_table, context_lens, page_size, scale=None, mask=None)
¶
Paged attention backed by the C++/Metal extension.
The Python wrapper keeps the model-facing shape as [B, H_q, L, D], while the extension sees flattened query heads and contiguous page storage.