Skip to main content

Mojo struct

MHATileSummary

@register_passable(trivial) 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, Movable, UnknownDestructibility

Aliases

__copyinit__is_trivial

alias __copyinit__is_trivial = ValidLengthType.__copyinit__is_trivial

__del__is_trivial

alias __del__is_trivial = ValidLengthType.__del__is_trivial

__moveinit__is_trivial

alias __moveinit__is_trivial = ValidLengthType.__moveinit__is_trivial

Methods

__init__

__init__(batch_size: UInt32, max_num_prompt_tiles: UInt32, valid_length: ValidLengthType, max_seq_len: UInt32) -> Self

get_current_work_info

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

Returns:

WorkInfo

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

Returns:

WorkInfo

unsafe_get_current_work_info

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

Returns:

WorkInfo

max_idx

max_idx(self, num_heads: UInt32) -> UInt32

Returns:

UInt32

grid_dim

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

Returns:

Tuple

seq_info

seq_info(self, work: WorkInfo) -> SeqInfo

Returns:

SeqInfo

unsafe_seq_info

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

Returns:

SeqInfo

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

Returns:

SeqInfo

Was this page helpful?