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_gpu_naive

def mha_gpu_naive[output_type: DType, k_t: MHAOperand, v_t: MHAOperand, mask_t: MHAMask, //, ragged: Bool = False, sink: Bool = False, _use_valid_length: Bool = False, _is_cache_length_accurate: Bool = False](q: LayoutTensor[element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment], k: k_t, v: v_t, mask_functor: mask_t, output: LayoutTensor[output_type, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], valid_length: LayoutTensor[DType.uint32, element_layout=valid_length.element_layout, layout_int_type=valid_length.layout_int_type, linear_idx_type=valid_length.linear_idx_type, masked=valid_length.masked, alignment=valid_length.alignment], scale: Float32, batch_size: Int, max_prompt_len: Int, max_cache_size: Int, num_heads: Int, depth: Int, group: Int, ctx: DeviceContext, sink_weights: OptionalReg[LayoutTensor[q.dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

Launch the naive (two-pass BMM) GPU attention implementation.

Computes attention as two separate batched matrix multiplications using temporary GMEM storage for the P = softmax(Q·Kᵀ / scale) intermediate. This is slower than flash attention but supports any head depth and serves as a correctness reference. Dispatches three sequential kernels: _bmm0_bs (Q·Kᵀ + mask), softmax normalisation, and _bmm1_bs (P·V).

Parameters:

output_type (DType): Element type of the attention output.
k_t (MHAOperand): Key operand type (dense or KV-cache).
v_t (MHAOperand): Value operand type (dense or KV-cache).
mask_t (MHAMask): Attention mask type.
ragged (Bool): True for ragged-batch inputs.
sink (Bool): True to enable attention-sink mode.
_use_valid_length (Bool): True to read per-sequence valid lengths.
_is_cache_length_accurate (Bool): True when cache length is exact.

Args:

q (LayoutTensor[element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment]): Query tensor.
k (k_t): Key operand.
v (v_t): Value operand.
mask_functor (mask_t): Mask instance.
output (LayoutTensor[output_type, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment]): Mutable output tensor.
valid_length (LayoutTensor[DType.uint32, element_layout=valid_length.element_layout, layout_int_type=valid_length.layout_int_type, linear_idx_type=valid_length.linear_idx_type, masked=valid_length.masked, alignment=valid_length.alignment]): Per-sequence valid lengths.
scale (Float32): Softmax temperature scale.
batch_size (Int): Number of sequences in the batch.
max_prompt_len (Int): Maximum query sequence length.
max_cache_size (Int): Maximum key/value sequence length.
num_heads (Int): Number of query heads.
depth (Int): Attention head depth.
group (Int): GQA group size.
ctx (DeviceContext): GPU device context.
sink_weights (OptionalReg[LayoutTensor[q.dtype, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Sink-token weights for attention-sink mode.

def mha_gpu_naive[q_type: DType, k_type: DType, v_type: DType, output_type: DType, mask_type: DType, //, sink: Bool = False](q: LayoutTensor[q_type, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment], k: LayoutTensor[k_type, element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment], v: LayoutTensor[v_type, element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment], mask: LayoutTensor[mask_type, element_layout=mask.element_layout, layout_int_type=mask.layout_int_type, linear_idx_type=mask.linear_idx_type, masked=mask.masked, alignment=mask.alignment], output: LayoutTensor[output_type, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], scale: Float32, batch_size: Int, seq_len: Int, num_keys: Int, num_heads: Int, depth: Int, group: Int, ctx: DeviceContext, sink_weights: OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

def mha_gpu_naive[q_type: DType, k_type: DType, v_type: DType, output_type: DType, MaskType: MHAMask, //, sink: Bool = False](q: LayoutTensor[q_type, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment], k: LayoutTensor[k_type, element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment], v: LayoutTensor[v_type, element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment], mask: MaskType, output: LayoutTensor[output_type, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], scale: Float32, batch_size: Int, seq_len: Int, num_keys: Int, num_heads: Int, depth: Int, group: Int, ctx: DeviceContext, sink_weights: OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

def mha_gpu_naive[q_type: DType, k_type: DType, v_type: DType, output_type: DType, mask_type: DType, q_tt_layout: TensorLayout, k_tt_layout: TensorLayout, v_tt_layout: TensorLayout, mask_tt_layout: TensorLayout, output_tt_layout: TensorLayout, //, sink: Bool = False](q: TileTensor[q_type, q_tt_layout, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: TileTensor[k_type, k_tt_layout, Storage=k.Storage, linear_idx_type=k.linear_idx_type], v: TileTensor[v_type, v_tt_layout, Storage=v.Storage, linear_idx_type=v.linear_idx_type], mask: TileTensor[mask_type, mask_tt_layout, Storage=mask.Storage, linear_idx_type=mask.linear_idx_type], output: TileTensor[output_type, output_tt_layout, Storage=output.Storage, linear_idx_type=output.linear_idx_type], scale: Float32, batch_size: Int, seq_len: Int, num_keys: Int, num_heads: Int, depth: Int, group: Int, ctx: DeviceContext, sink_weights: OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

TileTensor overload of mha_gpu_naive (materialized mask). Bridges to LayoutTensor internally.

Parameters:

q_type (DType): Element type of the query tensor.
k_type (DType): Element type of the key tensor.
v_type (DType): Element type of the value tensor.
output_type (DType): Element type of the output tensor.
mask_type (DType): Element type of the dense attention mask tensor.
q_tt_layout (TensorLayout): Compile-time TensorLayout of the query tensor.
k_tt_layout (TensorLayout): Compile-time TensorLayout of the key tensor.
v_tt_layout (TensorLayout): Compile-time TensorLayout of the value tensor.
mask_tt_layout (TensorLayout): Compile-time TensorLayout of the mask tensor.
output_tt_layout (TensorLayout): Compile-time TensorLayout of the output tensor.
sink (Bool): True to enable attention-sink mode (defaults to False).

Args:

q (TileTensor[q_type, q_tt_layout, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Query TileTensor with BSHD layout.
k (TileTensor[k_type, k_tt_layout, Storage=k.Storage, linear_idx_type=k.linear_idx_type]): Key TileTensor.
v (TileTensor[v_type, v_tt_layout, Storage=v.Storage, linear_idx_type=v.linear_idx_type]): Value TileTensor.
mask (TileTensor[mask_type, mask_tt_layout, Storage=mask.Storage, linear_idx_type=mask.linear_idx_type]): Dense attention mask TileTensor.
output (TileTensor[output_type, output_tt_layout, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Mutable output TileTensor.
scale (Float32): Softmax temperature scale applied to Q·Kᵀ.
batch_size (Int): Number of sequences in the batch.
seq_len (Int): Maximum query sequence length in the batch.
num_keys (Int): Maximum key/value sequence length.
num_heads (Int): Number of query heads.
depth (Int): Attention head depth (key/value dimension per head).
group (Int): GQA group size (query heads per key/value head).
ctx (DeviceContext): GPU device context for kernel dispatch.
sink_weights (OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink-token weight tensor for attention sinks.

def mha_gpu_naive[q_type: DType, k_type: DType, v_type: DType, output_type: DType, MaskType: MHAMask, q_tt_layout: TensorLayout, k_tt_layout: TensorLayout, v_tt_layout: TensorLayout, output_tt_layout: TensorLayout, //, sink: Bool = False](q: TileTensor[q_type, q_tt_layout, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: TileTensor[k_type, k_tt_layout, Storage=k.Storage, linear_idx_type=k.linear_idx_type], v: TileTensor[v_type, v_tt_layout, Storage=v.Storage, linear_idx_type=v.linear_idx_type], mask: MaskType, output: TileTensor[output_type, output_tt_layout, Storage=output.Storage, linear_idx_type=output.linear_idx_type], scale: Float32, batch_size: Int, seq_len: Int, num_keys: Int, num_heads: Int, depth: Int, group: Int, ctx: DeviceContext, sink_weights: OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)

TileTensor overload of mha_gpu_naive (MHAMask functor). Bridges to LayoutTensor internally.

Parameters:

q_type (DType): Element type of the query tensor.
k_type (DType): Element type of the key tensor.
v_type (DType): Element type of the value tensor.
output_type (DType): Element type of the output tensor.
MaskType (MHAMask): Attention mask type implementing MHAMask.
q_tt_layout (TensorLayout): Compile-time TensorLayout of the query tensor.
k_tt_layout (TensorLayout): Compile-time TensorLayout of the key tensor.
v_tt_layout (TensorLayout): Compile-time TensorLayout of the value tensor.
output_tt_layout (TensorLayout): Compile-time TensorLayout of the output tensor.
sink (Bool): True to enable attention-sink mode (defaults to False).

Args:

q (TileTensor[q_type, q_tt_layout, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Query TileTensor with BSHD layout.
k (TileTensor[k_type, k_tt_layout, Storage=k.Storage, linear_idx_type=k.linear_idx_type]): Key TileTensor.
v (TileTensor[v_type, v_tt_layout, Storage=v.Storage, linear_idx_type=v.linear_idx_type]): Value TileTensor.
mask (MaskType): Mask instance used to apply the attention mask.
output (TileTensor[output_type, output_tt_layout, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Mutable output TileTensor.
scale (Float32): Softmax temperature scale applied to Q·Kᵀ.
batch_size (Int): Number of sequences in the batch.
seq_len (Int): Maximum query sequence length in the batch.
num_keys (Int): Maximum key/value sequence length.
num_heads (Int): Number of query heads.
depth (Int): Attention head depth (key/value dimension per head).
group (Int): GQA group size (query heads per key/value head).
ctx (DeviceContext): GPU device context for kernel dispatch.
sink_weights (OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]]): Optional sink-token weight tensor for attention sinks.

def mha_gpu_naive[q_type: DType, output_type: DType, cache_t: KVCacheT, mask_t: MHAMask, //, ragged: Bool = False, sink: Bool = False](q: LayoutTensor[q_type, element_layout=q.element_layout, layout_int_type=q.layout_int_type, linear_idx_type=q.linear_idx_type, masked=q.masked, alignment=q.alignment], k: cache_t, v: cache_t, mask_functor: mask_t, output: LayoutTensor[output_type, element_layout=output.element_layout, layout_int_type=output.layout_int_type, linear_idx_type=output.linear_idx_type, masked=output.masked, alignment=output.alignment], valid_length: LayoutTensor[DType.uint32, element_layout=valid_length.element_layout, layout_int_type=valid_length.layout_int_type, linear_idx_type=valid_length.linear_idx_type, masked=valid_length.masked, alignment=valid_length.alignment], scale: Float32, batch_size: Int, max_prompt_len: Int, max_cache_size: Int, num_heads: Int, depth: Int, group: Int, ctx: DeviceContext, sink_weights: OptionalReg[LayoutTensor[q_type, Layout.row_major(Int(-1)), ImmutAnyOrigin]] = None)