IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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 struct

MHATileSummary

struct MHATileSummary[ValidLengthType: OptionalPointer]

Fields​

  • ​batch_size (UInt32):
  • ​max_num_prompt_tiles (UInt32):
  • ​valid_length (ValidLengthType):
  • ​max_seq_len (UInt32):

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, 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:

WorkInfo

def get_current_work_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, idx: MHATileState) -> WorkInfo

Returns:

WorkInfo

unsafe_get_current_work_info​

def unsafe_get_current_work_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, idx: UInt32) -> WorkInfo

Returns:

WorkInfo

max_idx​

def max_idx(self, num_heads: UInt32) -> UInt32

Returns:

UInt32

grid_dim​

static def grid_dim[num_heads: UInt32](max_num_prompt_tiles: UInt32, batch_size: UInt32) -> Tuple[Int, Int, Int]

Returns:

Tuple[Int, Int, Int]

seq_info​

def seq_info(self, work: WorkInfo) -> SeqInfo

Returns:

SeqInfo

unsafe_seq_info​

def unsafe_seq_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, idx: UInt32) -> SeqInfo

Returns:

SeqInfo

def unsafe_seq_info[tile_shape: UInt32, num_heads: UInt32, schedule: MHASchedule](self, state: MHATileState) -> SeqInfo

Returns:

SeqInfo