For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo struct

MHAAttentionConfigRDNA

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

Holds the fixed RDNA Wave32 multi-head attention configuration and index math.

Parameterizes a single attention kernel shape (16x16x16 WMMA) for both prefill and decode passes, and exposes the head, tile, and offset helpers the kernel uses to map block and lane IDs to Q/K/V rows in global memory.

Parameters

• ​token_gen (Bool): Whether the kernel is running in decode (token-by-token) generation mode versus prefill.
• ​config (MHAConfig[config.dtype]): The base multi-head attention configuration supplying head counts and tile dimensions.
• ​group (Int): Group-query attention group size (number of query heads sharing one key/value head).

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, 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 def q_head_idx() -> Int

Returns:

Int

q_tile_idx

static def q_tile_idx() -> Int

Returns:

Int

kv_head_idx

static def kv_head_idx() -> Int

Returns:

Int

get_mma_shape

static def get_mma_shape() -> IndexList[Int(3)]

Returns:

IndexList[Int(3)]

get_q_offset

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

Returns:

UInt32

get_output_offset

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

Returns:

UInt32

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit