Skip to main content

/

DocsCode

Mojo struct

MHAAttentionConfigRDNA

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

RDNA-specific attention configuration for Wave32 WMMA.

This config always uses:

  • MMA shape: 16x16x16 (only shape supported by RDNA WMMA)
  • k_group_size: 1
  • shared_kv: False (P reuses K's shared memory)
  • full_kv: False
  • depth_padded: True
  • double_buffer: False

Implemented traits​

AnyType, AttentionConfig, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable

comptime members​

depth_padded​

comptime depth_padded = True

double_buffer​

comptime double_buffer = False

full_kv​

comptime full_kv = False

shared_kv​

comptime shared_kv = False

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]

Get MMA shape - always 16x16x16 for RDNA.

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