Skip to main content

/

DocsCode

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 if MHAAttentionConfig[token_gen, config, group].USE_EXPERIMENTAL_CDNA4_MHA_KERNEL else True

double_buffer​

comptime double_buffer = True if MHAAttentionConfig[token_gen, config, group].USE_EXPERIMENTAL_CDNA4_MHA_KERNEL else False

full_kv​

comptime full_kv = True if MHAAttentionConfig[token_gen, config, group].USE_EXPERIMENTAL_CDNA4_MHA_KERNEL else False

shared_kv​

comptime shared_kv = False if MHAAttentionConfig[token_gen, config, group].USE_EXPERIMENTAL_CDNA4_MHA_KERNEL else True

USE_EXPERIMENTAL_CDNA4_MHA_KERNEL​

comptime USE_EXPERIMENTAL_CDNA4_MHA_KERNEL = _cdna_4_or_newer() and get_defined_bool["USE_EXPERIMENTAL_CDNA4_MHA_KERNEL", False]() and not token_gen

Methods​

q_head_idx​

static q_head_idx() -> UInt

Returns:

UInt

q_tile_idx​

static q_tile_idx() -> UInt

Returns:

UInt

kv_head_idx​

static kv_head_idx() -> UInt

Returns:

UInt

get_mma_shape​

static get_mma_shape() -> IndexList[3]

Returns:

IndexList

get_q_offset​

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

Returns:

UInt32

get_output_offset​

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

Returns:

UInt32

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact