Mojo struct

MHAConfig

@register_passable(trivial) struct MHAConfig[dtype: DType]

Fields

num_heads (UInt):
depth (UInt):
padded_depth (UInt):
num_queries_per_block (UInt):
num_keys_per_block (UInt):
BK (UInt):
WM (UInt):
WN (UInt):
num_pipeline_stages (UInt):
k_group_size (UInt):
algorithm (FlashAttentionAlgorithm):
swizzle_mode (TensorMapSwizzle):

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, Movable, UnknownDestructibility, Writable

Aliases

`copyinitis_trivial`

comptime __copyinit__is_trivial = True

`delis_trivial`

comptime __del__is_trivial = True

`moveinitis_trivial`

comptime __moveinit__is_trivial = True

Methods

`init`

__init__(num_heads: UInt, depth: UInt, num_queries_per_block: OptionalReg[UInt] = None, num_keys_per_block: OptionalReg[UInt] = None, BK: OptionalReg[UInt] = None, WM: OptionalReg[UInt] = None, WN: OptionalReg[UInt] = None, num_pipeline_stages: UInt = 4, k_group_size: UInt = 1, algorithm: FlashAttentionAlgorithm = FlashAttentionAlgorithm(-1), swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B) -> Self

`block_m`

block_m(self) -> UInt

Returns:

Fields​

Implemented traits​

Aliases​

__copyinit__is_trivial​

__del__is_trivial​

__moveinit__is_trivial​

Methods​

__init__​

block_m​

block_n​

block_k​

warp_m​

warp_n​

num_warps_m​

num_warps_n​

num_consumer_threads​

num_producer_threads​

num_threads​

swizzle_granularity​

q_smem_size​

kv_smem_size​

k_smem_size​

v_smem_size​

p_smem_size​

warp_scratch_smem_size​

shared_mem_bytes​

__str__​

write_to​

Fields

Implemented traits

Aliases

`copyinitis_trivial`

`delis_trivial`

`moveinitis_trivial`

Methods

`init`

`block_m`

`block_n`

`block_k`

`warp_m`

`warp_n`

`num_warps_m`

`num_warps_n`

`num_consumer_threads`

`num_producer_threads`

`num_threads`

`swizzle_granularity`

`q_smem_size`

`kv_smem_size`

`k_smem_size`

`v_smem_size`

`p_smem_size`

`warp_scratch_smem_size`

`shared_mem_bytes`

`str`

`write_to`