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_kernel

def fp8_index_kernel[dtype: DType, OutputLT: TensorLayout, QLT: TensorLayout, QSLT: TensorLayout, k_operand_type: MHAOperand, ks_operand_type: MHAOperand, block_tile_shape: Array[Int, Int(2)], VLLT: TensorLayout, num_heads: Int, depth: Int, _is_cache_length_accurate: Bool = False](output_tt: TileTensor[DType.float32, OutputLT, MutAnyOrigin], q_tt: TileTensor[dtype, QLT, ImmutAnyOrigin], q_s_tt: TileTensor[DType.float32, QSLT, MutAnyOrigin], k_operand: k_operand_type, ks_operand: ks_operand_type, valid_length_tt: TileTensor[DType.uint32, VLLT, ImmutAnyOrigin])

Computes the scalar FP8 index/gather score kernel as a Blackwell tensor-core fallback.

Each block computes a slice of the query sequence against a tile of the paged key cache, accumulating per-head logits in shared memory and writing the scale-weighted row sum to the output tensor.

Parameters:

dtype (DType): Data type of the query and key tiles.
OutputLT (TensorLayout): Layout type of the output score tensor.
QLT (TensorLayout): Layout type of the query tensor.
QSLT (TensorLayout): Layout type of the query scale tensor.
k_operand_type (MHAOperand): MHAOperand type used to address the paged key cache.
ks_operand_type (MHAOperand): MHAOperand type used to address the key scales.
block_tile_shape (Array[Int, Int(2)]): Block tile shape as [BM, BN] rows of sequence and keys.
VLLT (TensorLayout): Layout type of the valid-length (sequence offset) tensor.
num_heads (Int): Number of attention heads.
depth (Int): Per-head feature depth.
_is_cache_length_accurate (Bool): When True, cache_length already includes new tokens.

Args:

output_tt (TileTensor[DType.float32, OutputLT, MutAnyOrigin]): Output score tensor of shape [total_seq_len, num_keys].
q_tt (TileTensor[dtype, QLT, ImmutAnyOrigin]): Query tensor of shape [total_seq_len, num_heads, depth].
q_s_tt (TileTensor[DType.float32, QSLT, MutAnyOrigin]): Per-query scale tensor of shape [total_seq_len, num_heads].
k_operand (k_operand_type): Ragged paged operand providing key rows.
ks_operand (ks_operand_type): Ragged paged operand providing per-key scales.
valid_length_tt (TileTensor[DType.uint32, VLLT, ImmutAnyOrigin]): Cumulative sequence offsets of shape [batch_size + 1].