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

Mojo struct

AMDStructuredConfig

struct AMDStructuredConfig[config: MHAConfig[config.dtype], group: Int, token_gen: Bool = False, mla_mode: Bool = False, q_seq_len: Int = Int(1)]

Holds the tile layout and indexing helpers for GFX950 structured attention.

Wraps an MHAConfig with group-query and MLA metadata, deriving the comptime shared-memory decisions (shared_kv, full_kv, depth_padded, double_buffer) and exposing query/key-value head and tile index helpers used by the attention kernels.

Parameters

• ​config (MHAConfig[config.dtype]): The base multi-head attention configuration.
• ​group (Int): Number of query heads sharing one key-value head.
• ​token_gen (Bool): Selects decode (True) versus prefill (False) tiling.
• ​mla_mode (Bool): Enables multi-latent attention tiling when True.
• ​q_seq_len (Int): Token slots the decode fold stacks into M — the tile's height, padded above a sequence's real token count on narrow widths; selects the MFMA shape for the MHA fold's wide arm.

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, 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() <= Int(64))

full_kv

comptime full_kv = True

shared_kv

comptime shared_kv = token_gen and (config > Int(256))

Methods

heads_per_tile

static def heads_per_tile() -> Int

Returns:

Int

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