For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

mha_cross_gpu_naive

def mha_cross_gpu_naive[cache_t: KVCacheT, mask_t: MHAMask, dtype: DType, //, rank: Int](output: TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type], q_input_row_offsets: TileTensor[DType.uint32, Storage=q_input_row_offsets.Storage, address_space=q_input_row_offsets.address_space, linear_idx_type=q_input_row_offsets.linear_idx_type], q_max_seq_len: Int, k: cache_t, v: cache_t, kv_input_row_offsets: TileTensor[DType.uint32, Storage=kv_input_row_offsets.Storage, address_space=kv_input_row_offsets.address_space, linear_idx_type=kv_input_row_offsets.linear_idx_type], mask_functor: mask_t, scale: Float32, ctx: DeviceContext)

Naive cross attention on GPU.

Note that this assumes ragged tensor inputs and uses a mask functor.

Computes: (1) Transpose (Q) BSHD -> BHSD; (2) Transpose (K) BSHD -> BHSD; (3) Transpose (V) BSHD -> BHSD; (4) P = Bmm(Q, K), P is also called "score"; (5) P = P * scale + mask; (6) P = softmax(P); (7) O = Bmm(P, V) (8) Output = Transpose(O).

B, S, H, D denote batch size, sequence length, head count and depth, respectively. (1), (2), (3) happens while loading the data into shared memory. (8) happens when writing output to global memory.

All inputs (query, key, and value) must have BSHD layout. The mask can be BSS or BHSS.

This kernel also handles grouped attention optimization. In this case the shape of K and V are BShD where h = H / num_groups.

Parameters:

cache_t (KVCacheT): The paged KV cache type used for k and v (inferred).
mask_t (MHAMask): The mask functor type applied to attention scores (inferred).
dtype (DType): The element type of the query, key, value, and output tensors (inferred).
rank (Int): The number of dimensions of the input tensors. Must be 3 for ragged inputs.

Args:

output (TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type]): The output tensor receiving the cross attention result. Same dtype as q and the KV cache, in BSHD layout.
q (TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): The query tensor in BSHD layout. The static shape's last two dimensions give the head count and depth.
q_input_row_offsets (TileTensor[DType.uint32, Storage=q_input_row_offsets.Storage, address_space=q_input_row_offsets.address_space, linear_idx_type=q_input_row_offsets.linear_idx_type]): Per-batch start and end offsets into the ragged query tensor. Length is batch_size + 1.
q_max_seq_len (Int): The maximum query sequence length across the batch.
k (cache_t): The paged KV cache holding keys.
v (cache_t): The paged KV cache holding values.
kv_input_row_offsets (TileTensor[DType.uint32, Storage=kv_input_row_offsets.Storage, address_space=kv_input_row_offsets.address_space, linear_idx_type=kv_input_row_offsets.linear_idx_type]): Per-batch start and end offsets into the ragged KV input. Length is batch_size + 1.
mask_functor (mask_t): The mask instance applied to attention scores before softmax.
scale (Float32): The scaling factor multiplied with the query-key scores.
ctx (DeviceContext): The device context used to enqueue GPU kernels and buffers.