Mojo struct

MHAAttentionConfig

struct MHAAttentionConfig[token_gen: Bool, config: MHAConfig[config.dtype], group: Int]

Implemented traits

AnyType, AttentionConfig, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable

`comptime` members

`depth_padded`

comptime depth_padded = False

`double_buffer`

comptime double_buffer = not token_gen

`double_buffer_k_only`

comptime double_buffer_k_only = token_gen and (config.block_n() <= 64)

`full_kv`

comptime full_kv = True

`shared_kv`

comptime shared_kv = token_gen and (config > 256)

Methods

`q_head_idx`

static q_head_idx() -> Int

Returns:

Int

`q_tile_idx`

static q_tile_idx() -> Int

Returns:

Int

`kv_head_idx`

static kv_head_idx() -> Int

Returns:

Int

`get_mma_shape`

static get_mma_shape() -> IndexList[3]

Returns:

IndexList

`get_q_offset`

static get_q_offset[q_depth: Int]() -> UInt32

Returns:

UInt32

`get_output_offset`

static get_output_offset[output_depth: Int]() -> UInt32

Returns:

UInt32

Implemented traits​

comptime members​

depth_padded​

double_buffer​

double_buffer_k_only​

full_kv​

shared_kv​

Methods​

q_head_idx​

q_tile_idx​

kv_head_idx​

get_mma_shape​

get_q_offset​

get_output_offset​

Implemented traits

`comptime` members

`depth_padded`

`double_buffer`

`double_buffer_k_only`

`full_kv`

`shared_kv`

Methods

`q_head_idx`

`q_tile_idx`

`kv_head_idx`

`get_mma_shape`

`get_q_offset`

`get_output_offset`