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
MHATileSummary
struct MHATileSummary[ValidLengthType: OptionalPointer]
Summarizes the dimensions and valid-length metadata of the MHA work grid.
Stores the batch size, maximum number of prompt tiles, optional per-batch
sequence length offsets, and the maximum sequence length, providing the
mapping from linear work-tile indices to (prompt_tile, head, batch)
coordinates.
Parameters
- ValidLengthType (
OptionalPointer): The optional pointer type carrying per-batch sequence length offsets (inferred).
Fields
- batch_size (
UInt32): - max_num_prompt_tiles (
UInt32): - valid_length (
ValidLengthType): - max_seq_len (
UInt32):
Implemented traits
AnyType,
Copyable,
Deinitable,
ImplicitlyCopyable,
Movable,
RegisterPassable,
TrivialRegisterPassable
Methods
__init__
def __init__(batch_size: UInt32, max_num_prompt_tiles: UInt32, valid_length: ValidLengthType, max_seq_len: UInt32) -> Self
get_current_work_info
def get_current_work_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, idx: UInt32) -> WorkInfo
Returns:
def get_current_work_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, idx: MHATileState) -> WorkInfo
Returns:
unsafe_get_current_work_info
def unsafe_get_current_work_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, idx: UInt32) -> WorkInfo
Returns:
max_idx
grid_dim
static def grid_dim[num_heads: UInt32](max_num_prompt_tiles: UInt32, batch_size: UInt32) -> Tuple[Int, Int, Int]
Returns:
seq_info
unsafe_seq_info
def unsafe_seq_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, idx: UInt32) -> SeqInfo
Returns:
def unsafe_seq_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, state: MHATileState) -> SeqInfo
Returns: