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
MHAConfig
struct MHAConfig[dtype: DType]
Compile-time and runtime tile-shape configuration for MHA GPU kernels.
Stores the tile dimensions (BM, BN, BK), warp tile dimensions (WM, WN),
pipeline depth, and algorithm variant used when launching flash-attention
kernels. The constructor auto-selects sensible defaults based on dtype
and the detected GPU architecture when optional fields are left as None.
Parameters
- dtype (
DType): The element data type shared by Q, K, V, and the output tensor.
Fields
- num_heads (
Int): - depth (
Int): - padded_depth (
Int): - num_queries_per_block (
Int): - num_keys_per_block (
Int): - BK (
Int): - WM (
Int): - WN (
Int): - num_pipeline_stages (
Int): - k_group_size (
Int): - algorithm (
FlashAttentionAlgorithm): - swizzle_mode (
TensorMapSwizzle):
Implemented traits
AnyType,
Copyable,
Deinitable,
ImplicitlyCopyable,
Movable,
RegisterPassable,
TrivialRegisterPassable,
Writable
Methods
__init__
def __init__(num_heads: Int, depth: Int, num_queries_per_block: Optional[Int] = None, num_keys_per_block: Optional[Int] = None, BK: Optional[Int] = None, WM: Optional[Int] = None, WN: Optional[Int] = None, num_pipeline_stages: Int = Int(4), k_group_size: Int = Int(1), algorithm: FlashAttentionAlgorithm = FlashAttentionAlgorithm(Int(-1)), swizzle_mode: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B) -> 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
swizzle_granularity
q_smem_size
kv_smem_size
k_smem_size
v_smem_size
p_smem_size
warp_scratch_smem_size
shared_mem_bytes
write_to
def write_to(self, mut writer: T)