Skip to main content

Mojo struct

MHAAttentionConfigRDNA

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

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable

comptime members​

depth_padded​

comptime depth_padded = True

double_buffer​

comptime double_buffer = False

double_buffer_k_only​

comptime double_buffer_k_only = False

full_kv​

comptime full_kv = False

shared_kv​

comptime shared_kv = False

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[3]

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