Skip to main content

Mojo trait

MHATileScheduler

Implemented traits

AnyType, Copyable, UnknownDestructibility

Aliases

__copyinit__is_trivial

alias __copyinit__is_trivial

A flag (often compiler generated) to indicate whether the implementation of __copyinit__ is trivial.

The implementation of __copyinit__ is considered to be trivial if:

  • The struct has a compiler-generated trivial __copyinit__ and all its fields have a trivial __copyinit__ method.

In practice, it means the value can be copied by copying the bits from one location to another without side effects.

__del__is_trivial

alias __del__is_trivial

A flag (often compiler generated) to indicate whether the implementation of __del__ is trivial.

The implementation of __del__ is considered to be trivial if:

  • The struct has a compiler-generated trivial destructor and all its fields have a trivial __del__ method.

In practice, it means that the __del__ can be considered as no-op.

may_advance

alias may_advance

mha_schedule

alias mha_schedule

The MHATileScheduler trait describes a schedule for the persistent kernel.

Required methods

__copyinit__

__copyinit__(out self: _Self, existing: _Self, /)

Create a new instance of the value by copying an existing one.

Args:

  • existing (_Self): The value to copy.

Returns:

_Self

get_current_work_info

get_current_work_info[ValidLengthType: OptionalPointer, //](self: _Self, ts: MHATileSummary[ValidLengthType], state: MHATileState) -> WorkInfo

Returns the current WorkInfo.

Returns:

WorkInfo

advance

advance[ValidLengthType: OptionalPointer, //, producer: Bool, sync: MHASchedulerSynchronization = MHASchedulerSynchronization(1)](self: _Self, ts: MHATileSummary[ValidLengthType], mut state: MHATileState, pipeline_idx: UInt32) -> OptionalReg[SeqInfo]

Advance state to the next work item. func must return a Bool indicating whether there is more work. Returns True if there is more work.

Returns:

OptionalReg

grid_dim

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

Return the grid_dim required for the kernel.

Returns:

Tuple

initial_state

initial_state[ValidLengthType: OptionalPointer, //](self: _Self, ptr: UnsafePointer[UInt32, address_space=AddressSpace(3)], tile_summary: MHATileSummary[ValidLengthType]) -> MHATileState

Create the initial state object.

Returns:

MHATileState

unsafe_seq_info

unsafe_seq_info[ValidLengthType: OptionalPointer, //](self: _Self, ts: MHATileSummary[ValidLengthType], state: MHATileState) -> SeqInfo

Returns:

SeqInfo

Provided methods

copy

copy(self: _Self) -> _Self

Explicitly construct a copy of self.

Returns:

_Self: A copy of this value.

Was this page helpful?