For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Mojo module
mha_utils
Shared configuration types and dispatch helpers for MHA GPU kernels.
Provides MHAConfig (tile/pipeline configuration), FlashAttentionAlgorithm
(algorithm variant selector), mask-dispatch helpers, and partition-scheme
types used by both prefill and decode attention kernels.
comptime values
callback_fn_type
comptime callback_fn_type = def[mask_t: MHAMask](mask: mask_t) raises capturing thin -> None
is_sm100
comptime is_sm100 = (StringSlice("sm_100") in String(_accelerator_arch())) or (StringSlice("sm_103") in String(_accelerator_arch()))
is_sm90
comptime is_sm90 = (StringSlice("sm_90") in String(_accelerator_arch()))
is_sm90or100
comptime is_sm90or100 = is_sm90 or is_sm100
MHA_PDL_LEVEL
comptime MHA_PDL_LEVEL = PDLLevel.OVERLAP_AT_END if get_defined_bool[StringSlice("MHA_PDL"), True]() else PDLLevel.OFF
Structs
-
DynamicInt: A runtime integer value that satisfiesOptionallyStaticInt. -
FlashAttentionAlgorithm: Identifies which flash-attention algorithm variant to use for a kernel launch. -
MHAConfig: Compile-time and runtime tile-shape configuration for MHA GPU kernels. -
NoPartition: A single-partition (non-split-K) scheme for MHA decoding. -
SplitKPartition: A multi-partition split-K scheme for MHA decoding over long sequences. -
StaticInt: A compile-time constant integer that satisfiesOptionallyStaticInt.
Traits
-
MHAPartitionScheme: Trait describing how the key-value sequence is partitioned for split-K decoding. -
OptionallyStaticInt: Trait for integer values that may be statically known at compile time.
Functions
-
as_dynamic_row_major_1d: Reinterprets a generic-addressLayoutTensoras a 1-D dynamic row-major tensor. -
dispatch_mask: Instantiate anMHAMaskby name and invoke a callback with it. -
dispatch_materialized_mask: Wrap a dense mask tensor in aMaterializedMaskand invoke a callback. -
get_start_and_end_for_partitions: Calculate start and end indices for a partition.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!