For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /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_lengthalready 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].