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

generic_flash_attention_kv_cache_ragged_rel_logits

def generic_flash_attention_kv_cache_ragged_rel_logits[collection_t: KVCollectionT, dtype: DType, //, *, target: StringSpan[ImmStaticOrigin], local_window_size: Int = Int(-1), output_dtype: DType = dtype](q: LayoutTensor[dtype, 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], input_row_offsets: LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin], kv_collection: collection_t, layer_idx: UInt32, scale: Float32, bias: LayoutTensor[dtype, element_layout=bias.element_layout, layout_int_type=bias.layout_int_type, linear_idx_type=bias.linear_idx_type, masked=bias.masked, alignment=bias.alignment], cache_lengths: LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), ImmutAnyOrigin], output: LayoutTensor[output_dtype, 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], context: DeviceContext, decode_dispatch_metadata: MHADecodeDispatchMetadata)

Flash attention over a ragged batch with a relative-position bias.

bias is gathered by rel_dist = q_pos - k_pos inside the kernel via RelativeLogitsMask; local_window_size picks the visibility mask (-1 -> global causal). cache_lengths is the same tensor used to build kv_collection, passed again so the mask can recover each query's ragged-flat row into bias (the collection only exposes a per-sequence scalar accessor, not the raw tensor).