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

fp8_index_naive

def fp8_index_naive[dtype: DType, //, num_heads: Int, depth: Int](output: TileTensor[DType.float32, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], q: TileTensor[dtype, Storage=q.Storage, address_space=q.address_space, linear_idx_type=q.linear_idx_type], q_s: TileTensor[DType.float32, Storage=q_s.Storage, address_space=q_s.address_space, linear_idx_type=q_s.linear_idx_type], k: TileTensor[dtype, Storage=k.Storage, address_space=k.address_space, linear_idx_type=k.linear_idx_type], k_s: TileTensor[DType.float32, Storage=k_s.Storage, address_space=k_s.address_space, linear_idx_type=k_s.linear_idx_type], valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, address_space=valid_length.address_space, linear_idx_type=valid_length.linear_idx_type], cache_row_offsets: TileTensor[DType.uint32, Storage=cache_row_offsets.Storage, address_space=cache_row_offsets.address_space, linear_idx_type=cache_row_offsets.linear_idx_type], batch_size: Int, max_seq_len: Int, max_num_keys: Int, ctx: DeviceContext)

Computes the FP8 index/gather score via a two-pass matmul-then-reduce reference path.

Enqueues _index_matmul_max to produce per-head logits followed by _reduce_logits to sum across heads and apply the per-key scale, serving as a correctness reference for the optimized tensor-core kernels.

Parameters:

dtype (DType): Data type of the query and key tensors.
num_heads (Int): Number of attention heads.
depth (Int): Per-head feature depth.

Args:

output (TileTensor[DType.float32, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Output score tensor of shape [total_seq_len, max_num_keys].
q (TileTensor[dtype, Storage=q.Storage, address_space=q.address_space, linear_idx_type=q.linear_idx_type]): Query tensor of shape [total_seq_len, num_heads, depth].
q_s (TileTensor[DType.float32, Storage=q_s.Storage, address_space=q_s.address_space, linear_idx_type=q_s.linear_idx_type]): Per-query scale tensor of shape [total_seq_len, num_heads].
k (TileTensor[dtype, Storage=k.Storage, address_space=k.address_space, linear_idx_type=k.linear_idx_type]): Key tensor of shape [total_keys, 1, depth].
k_s (TileTensor[DType.float32, Storage=k_s.Storage, address_space=k_s.address_space, linear_idx_type=k_s.linear_idx_type]): Per-key scale tensor of shape [total_keys].
valid_length (TileTensor[DType.uint32, Storage=valid_length.Storage, address_space=valid_length.address_space, linear_idx_type=valid_length.linear_idx_type]): Cumulative sequence offsets of shape [batch_size + 1].
cache_row_offsets (TileTensor[DType.uint32, Storage=cache_row_offsets.Storage, address_space=cache_row_offsets.address_space, linear_idx_type=cache_row_offsets.linear_idx_type]): Per-batch row offsets into the paged key cache.
batch_size (Int): Number of sequences in the batch.
max_seq_len (Int): Maximum sequence length across the batch.
max_num_keys (Int): Maximum key count across the batch.
ctx (DeviceContext): Device context used to enqueue the kernels.

Raises:

When the underlying kernel enqueue reports a device-side error.