Skip to main content

/

DocsCode

Mojo struct

MHAAttentionConfigRDNA

struct MHAAttentionConfigRDNA[dtype: DType, //, token_gen: Bool, config: MHAConfig[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​

__copyinit__is_trivial​

comptime __copyinit__is_trivial = True

__del__is_trivial​

comptime __del__is_trivial = True

__moveinit__is_trivial​

comptime __moveinit__is_trivial = True

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: Scalar[DType.uint]]() -> UInt32

Returns:

UInt32

get_output_offset​

static get_output_offset[output_depth: Scalar[DType.uint]]() -> UInt32

Returns:

UInt32

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact