Mojo struct
MHAConfig
@register_passable(trivial)
struct MHAConfig
Fields
- type (
DType): - num_heads (
UInt): - 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):
Implemented traits
AnyType,
Copyable,
Movable,
UnknownDestructibility,
Writable
Methods
__init__
__init__(type: DType, num_heads: UInt, depth: UInt, num_queries_per_block: OptionalReg[UInt] = OptionalReg[UInt]({:i1 0, 1}), num_keys_per_block: OptionalReg[UInt] = OptionalReg[UInt]({:i1 0, 1}), BK: OptionalReg[UInt] = OptionalReg[UInt]({:i1 0, 1}), WM: OptionalReg[UInt] = OptionalReg[UInt]({:i1 0, 1}), WN: OptionalReg[UInt] = OptionalReg[UInt]({:i1 0, 1}), num_pipeline_stages: UInt = UInt(4), k_group_size: UInt = UInt(1), algorithm: FlashAttentionAlgorithm = FlashAttentionAlgorithm(-1)) -> Self
block_m
block_n
block_k
warp_m
warp_n
num_warps_m
num_warps_n
num_consumer_threads
num_producer_threads
num_threads
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
write_to[W: Writer](self, mut writer: W)
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!