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 struct
MHAPosition
struct MHAPosition[BM: Int, BN: Int, depth: Int, padded_depth: Int, q_num_heads: Int, group: Int, decoding: Bool]
Position of the MHA-kernel. When decoding=False, q_head_stride == q_num_heads. When decoding=True, q_head_stride == 1.
Parameters
- BM (
Int): Tile block size in the query (row) dimension, in elements. - BN (
Int): Tile block size in the key (column) dimension, in elements. - depth (
Int): Head dimension of the attention layer, in elements. - padded_depth (
Int): Head dimension padded to the tensor-core alignment boundary, in elements. - q_num_heads (
Int): Number of query attention heads. - group (
Int): Grouped-query attention group size, in query heads per KV head. - decoding (
Bool): Whether the kernel runs in single-token decoding mode.
Fields
- q_row (
UInt32): - q_col (
UInt32): - q_out_offset (
Int): - num_keys (
UInt32): - start_pos (
UInt32): - seq_len (
UInt32): - head_idx (
UInt32): - prompt_offset (
UInt32): - prompt_idx (
UInt32):
Implemented traits
AnyType,
Copyable,
Deinitable,
ImplicitlyCopyable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
num_q_heads_per_thread
comptime num_q_heads_per_thread = min(Int(2), ceildiv(group, Int(8))) if decoding else Int(1)
q_output_gmem_layout
comptime q_output_gmem_layout = Layout(IntTuple(BM, depth), IntTuple(depth if decoding else Int((mul depth, q_num_heads)), Int(1)))
q_stride
comptime q_stride = depth if decoding else (depth * q_num_heads)
split_gmem_layout
comptime split_gmem_layout = Layout(IntTuple((BM // Int(2)), depth), IntTuple(depth if decoding else Int((mul depth, q_num_heads)), Int(1)))
Methods
__init__
def __init__(q_row: UInt32, q_col: UInt32, q_out_offset: Int, num_keys: UInt32, start_pos: UInt32, seq_info: SeqInfo) -> Self
__eq__
__ne__
q_head_idx
kv_head_idx
write_to
def write_to(self, mut writer: T)
q_tile_num_rows
q_out_gmem_tensor
def q_out_gmem_tensor[dtype: DType](self, ptr: Pointer[Scalar[dtype]]) -> LayoutTensor[dtype, Self.q_output_gmem_layout, ptr.origin, layout_int_type=DType.int32, linear_idx_type=DType.int32, masked=True]
Returns:
mask_status
def mask_status[MaskType: MHAMask](self, mask: MaskType, kv_tile_start_row: UInt32) -> TileMaskStatus
Returns:
get_score_row
exp_sum_qk_max_ptr
def exp_sum_qk_max_ptr[partition_t: MHAPartitionScheme](self, partition: partition_t, batch_size: UInt32) -> Tuple[Pointer[Scalar[partition_t.accum_dtype], MutAnyOrigin], Pointer[Scalar[partition_t.accum_dtype], MutAnyOrigin]]
Returns:
get_start_and_end_for_partitions
def get_start_and_end_for_partitions[PartitionType: MHAPartitionScheme, MaskType: MHAMask, //, page_size: Int](self, partition: PartitionType, mask: MaskType) -> Tuple[UInt32, UInt32]
Returns:
get_q_gmem_row
static def get_q_gmem_row[MaxSeqLenType: OptionallyStaticInt, //, ragged: Bool](seq_info: SeqInfo, max_seq_len: MaxSeqLenType) -> UInt32
Returns:
static def get_q_gmem_row[ragged: Bool](seq_info: SeqInfo, max_seq_len: UInt32) -> UInt32
Returns: